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
606impl From<ColumnName> for mz_sql_parser::ast::Ident {
607    fn from(value: ColumnName) -> Self {
608        // Note: ColumnNames are known to be less than the max length of an Ident (I think?).
609        mz_sql_parser::ast::Ident::new_unchecked(value.0)
610    }
611}
612
613#[cfg(any(test, feature = "proptest"))]
614impl proptest::arbitrary::Arbitrary for ColumnName {
615    type Parameters = ();
616    type Strategy = BoxedStrategy<ColumnName>;
617
618    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
619        // Long column names are generally uninteresting, and can greatly
620        // increase the runtime for a test case, so bound the max length.
621        let mut weights = vec![(50, Just(1..8)), (20, Just(8..16))];
622        if std::env::var("PROPTEST_LARGE_DATA").is_ok() {
623            weights.extend([
624                (5, Just(16..128)),
625                (1, Just(128..1024)),
626                (1, Just(1024..4096)),
627            ]);
628        }
629        let name_length = Union::new_weighted(weights);
630
631        // Non-ASCII characters are also generally uninteresting and can make
632        // debugging harder.
633        let char_strat = Rc::new(Union::new_weighted(vec![
634            (50, proptest::char::range('A', 'z').boxed()),
635            (1, any::<char>().boxed()),
636        ]));
637
638        name_length
639            .prop_flat_map(move |length| proptest::collection::vec(Rc::clone(&char_strat), length))
640            .prop_map(|chars| ColumnName(chars.into_iter().collect::<Box<str>>()))
641            .no_shrink()
642            .boxed()
643    }
644}
645
646/// Default name of a column (when no other information is known).
647pub const UNKNOWN_COLUMN_NAME: &str = "?column?";
648
649/// Stable index of a column in a [`RelationDesc`].
650#[derive(
651    Clone,
652    Copy,
653    Debug,
654    Eq,
655    PartialEq,
656    PartialOrd,
657    Ord,
658    Serialize,
659    Deserialize,
660    Hash
661)]
662pub struct ColumnIndex(usize);
663
664#[cfg(any(test, feature = "proptest"))]
665static_assertions::assert_not_impl_all!(ColumnIndex: Arbitrary);
666
667impl ColumnIndex {
668    /// Returns a stable identifier for this [`ColumnIndex`].
669    pub fn to_stable_name(&self) -> String {
670        self.0.to_string()
671    }
672
673    pub fn to_raw(&self) -> usize {
674        self.0
675    }
676
677    pub fn from_raw(val: usize) -> Self {
678        ColumnIndex(val)
679    }
680}
681
682/// The version a given column was added at.
683#[derive(
684    Clone,
685    Copy,
686    Debug,
687    Eq,
688    PartialEq,
689    PartialOrd,
690    Ord,
691    Serialize,
692    Deserialize,
693    Hash
694)]
695#[cfg_attr(any(test, feature = "proptest"), derive(Arbitrary))]
696pub struct RelationVersion(u64);
697
698impl RelationVersion {
699    /// Returns the "root" or "initial" version of a [`RelationDesc`].
700    pub fn root() -> Self {
701        RelationVersion(0)
702    }
703
704    /// Returns an instance of [`RelationVersion`] which is "one" higher than `self`.
705    pub fn bump(&self) -> Self {
706        let next_version = self
707            .0
708            .checked_add(1)
709            .expect("added more than u64::MAX columns?");
710        RelationVersion(next_version)
711    }
712
713    /// Consume a [`RelationVersion`] returning the raw value.
714    ///
715    /// Should __only__ be used for serialization.
716    pub fn into_raw(self) -> u64 {
717        self.0
718    }
719
720    /// Create a [`RelationVersion`] from a raw value.
721    ///
722    /// Should __only__ be used for serialization.
723    pub fn from_raw(val: u64) -> RelationVersion {
724        RelationVersion(val)
725    }
726}
727
728impl From<RelationVersion> for SchemaId {
729    fn from(value: RelationVersion) -> Self {
730        SchemaId(usize::cast_from(value.0))
731    }
732}
733
734impl From<mz_sql_parser::ast::Version> for RelationVersion {
735    fn from(value: mz_sql_parser::ast::Version) -> Self {
736        RelationVersion(value.into_inner())
737    }
738}
739
740impl fmt::Display for RelationVersion {
741    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
742        write!(f, "v{}", self.0)
743    }
744}
745
746impl From<RelationVersion> for mz_sql_parser::ast::Version {
747    fn from(value: RelationVersion) -> Self {
748        mz_sql_parser::ast::Version::new(value.0)
749    }
750}
751
752impl RustType<ProtoRelationVersion> for RelationVersion {
753    fn into_proto(&self) -> ProtoRelationVersion {
754        ProtoRelationVersion { value: self.0 }
755    }
756
757    fn from_proto(proto: ProtoRelationVersion) -> Result<Self, TryFromProtoError> {
758        Ok(RelationVersion(proto.value))
759    }
760}
761
762/// Semantic type annotation for a column in a builtin catalog relation.
763///
764/// These are compile-time metadata used by the catalog ontology layer to
765/// describe the meaning of a column (e.g., that it contains a catalog item ID
766/// or a role ID). Possible values correspond to the entries in
767/// `SEMANTIC_TYPE_DEFS` in the `mz-catalog` crate.
768#[derive(
769    Clone,
770    Copy,
771    Debug,
772    PartialEq,
773    Eq,
774    PartialOrd,
775    Ord,
776    Hash,
777    serde::Serialize
778)]
779pub enum SemanticType {
780    CatalogItemId,
781    GlobalId,
782    ClusterId,
783    ReplicaId,
784    SchemaId,
785    DatabaseId,
786    RoleId,
787    NetworkPolicyId,
788    ShardId,
789    OID,
790    ObjectType,
791    ConnectionType,
792    SourceType,
793    MzTimestamp,
794    WallclockTimestamp,
795    ByteCount,
796    RecordCount,
797    CreditRate,
798    SqlDefinition,
799    RedactedSqlDefinition,
800}
801
802impl fmt::Display for SemanticType {
803    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
804        let s = match self {
805            SemanticType::CatalogItemId => "CatalogItemId",
806            SemanticType::GlobalId => "GlobalId",
807            SemanticType::ClusterId => "ClusterId",
808            SemanticType::ReplicaId => "ReplicaId",
809            SemanticType::SchemaId => "SchemaId",
810            SemanticType::DatabaseId => "DatabaseId",
811            SemanticType::RoleId => "RoleId",
812            SemanticType::NetworkPolicyId => "NetworkPolicyId",
813            SemanticType::ShardId => "ShardId",
814            SemanticType::OID => "OID",
815            SemanticType::ObjectType => "ObjectType",
816            SemanticType::ConnectionType => "ConnectionType",
817            SemanticType::SourceType => "SourceType",
818            SemanticType::MzTimestamp => "MzTimestamp",
819            SemanticType::WallclockTimestamp => "WallclockTimestamp",
820            SemanticType::ByteCount => "ByteCount",
821            SemanticType::RecordCount => "RecordCount",
822            SemanticType::CreditRate => "CreditRate",
823            SemanticType::SqlDefinition => "SqlDefinition",
824            SemanticType::RedactedSqlDefinition => "RedactedSqlDefinition",
825        };
826        f.write_str(s)
827    }
828}
829
830/// Metadata (other than type) for a column in a [`RelationDesc`].
831#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash)]
832struct ColumnMetadata {
833    /// Name of the column.
834    name: ColumnName,
835    /// Index into a [`SqlRelationType`] for this column.
836    typ_idx: usize,
837    /// Version this column was added at.
838    added: RelationVersion,
839    /// Version this column was dropped at.
840    dropped: Option<RelationVersion>,
841}
842
843/// A description of the shape of a relation.
844///
845/// It bundles a [`SqlRelationType`] with `ColumnMetadata` for each column in
846/// the relation.
847///
848/// # Examples
849///
850/// A `RelationDesc`s is typically constructed via its builder API:
851///
852/// ```
853/// use mz_repr::{SqlColumnType, RelationDesc, SqlScalarType};
854///
855/// let desc = RelationDesc::builder()
856///     .with_column("id", SqlScalarType::Int64.nullable(false))
857///     .with_column("price", SqlScalarType::Float64.nullable(true))
858///     .finish();
859/// ```
860///
861/// In more complicated cases, like when constructing a `RelationDesc` in
862/// response to user input, it may be more convenient to construct a relation
863/// type first, and imbue it with column names to form a `RelationDesc` later:
864///
865/// ```
866/// use mz_repr::RelationDesc;
867///
868/// # fn plan_query(_: &str) -> mz_repr::SqlRelationType { mz_repr::SqlRelationType::new(vec![]) }
869/// let relation_type = plan_query("SELECT * FROM table");
870/// let names = (0..relation_type.arity()).map(|i| match i {
871///     0 => "first",
872///     1 => "second",
873///     _ => "unknown",
874/// });
875/// let desc = RelationDesc::new(relation_type, names);
876/// ```
877///
878/// Next to the [`SqlRelationType`] we maintain a map of `ColumnIndex` to
879/// `ColumnMetadata`, where [`ColumnIndex`] is a stable identifier for a
880/// column throughout the lifetime of the relation. This allows a
881/// [`RelationDesc`] to represent a projection over a version of itself.
882///
883/// ```
884/// use std::collections::BTreeSet;
885/// use mz_repr::{ColumnIndex, RelationDesc, SqlScalarType};
886///
887/// let desc = RelationDesc::builder()
888///     .with_column("name", SqlScalarType::String.nullable(false))
889///     .with_column("email", SqlScalarType::String.nullable(false))
890///     .finish();
891///
892/// // Project away the second column.
893/// let demands = BTreeSet::from([1]);
894/// let proj = desc.apply_demand(&demands);
895///
896/// // We projected away the first column.
897/// assert!(!proj.contains_index(&ColumnIndex::from_raw(0)));
898/// // But retained the second.
899/// assert!(proj.contains_index(&ColumnIndex::from_raw(1)));
900///
901/// // The underlying `SqlRelationType` also contains a single column.
902/// assert_eq!(proj.typ().arity(), 1);
903/// ```
904///
905/// To maintain this stable mapping and track the lifetime of a column (e.g.
906/// when adding or dropping a column) we use `ColumnMetadata`. It maintains
907/// the index in [`SqlRelationType`] that corresponds to a given column, and the
908/// version at which this column was added or dropped.
909///
910#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
911pub struct RelationDesc {
912    typ: SqlRelationType,
913    metadata: BTreeMap<ColumnIndex, ColumnMetadata>,
914}
915
916impl RustType<ProtoRelationDesc> for RelationDesc {
917    // NOTE: `ProtoRelationDesc` has no field for the `ColumnIndex` keys, only the values in
918    // `ColumnIndex` order, so a desc whose indexes are sparse (what
919    // `VersionedRelationDesc::at_version` returns once a column has been dropped, and what
920    // `RelationDesc::apply_demand` returns) comes back from `from_proto` renumbered to `0..n`.
921    // `ColumnIndex::to_stable_name` is the arrow field name that `RowColumnarEncoder` and
922    // `RowColumnarDecoder` agree on, so encoding a sparse desc as the schema of data written
923    // under the original indexes builds a decoder that looks up the wrong fields. Only dense
924    // descs may be handed to `Codec::encode_schema`.
925    fn into_proto(&self) -> ProtoRelationDesc {
926        let (names, metadata): (Vec<_>, Vec<_>) = self
927            .metadata
928            .values()
929            .map(|meta| {
930                let metadata = ProtoColumnMetadata {
931                    added: Some(meta.added.into_proto()),
932                    dropped: meta.dropped.map(|v| v.into_proto()),
933                };
934                (meta.name.into_proto(), metadata)
935            })
936            .unzip();
937
938        // `metadata` Migration Logic: We wrote some `ProtoRelationDesc`s into Persist before the
939        // metadata field was added. To make sure our serialization roundtrips the same as before
940        // we added the field, we omit `metadata` if all of the values are equal to the default.
941        //
942        // Note: This logic needs to exist approximately forever.
943        let is_all_default_metadata = metadata.iter().all(|meta| {
944            meta.added == Some(RelationVersion::root().into_proto()) && meta.dropped == None
945        });
946        let metadata = if is_all_default_metadata {
947            Vec::new()
948        } else {
949            metadata
950        };
951
952        ProtoRelationDesc {
953            typ: Some(self.typ.into_proto()),
954            names,
955            metadata,
956        }
957    }
958
959    fn from_proto(proto: ProtoRelationDesc) -> Result<Self, TryFromProtoError> {
960        let typ: SqlRelationType = proto.typ.into_rust_if_some("ProtoRelationDesc::typ")?;
961
962        // Reject shapes that `VersionedRelationDesc::validate` calls corruption. Nothing
963        // downstream catches them: they decode and re-encode cleanly, and only panic at first
964        // use, e.g. `iter()` indexing `typ.columns()[typ_idx]` out of bounds, or `into_iter()`
965        // tripping `zip_eq`. Both are reachable from untrusted proto bytes.
966        if proto.names.len() != typ.column_types.len() {
967            return Err(TryFromProtoError::InvalidFieldError(format!(
968                "ProtoRelationDesc: names ({}) and column_types ({}) length mismatch",
969                proto.names.len(),
970                typ.column_types.len()
971            )));
972        }
973        if let Some(key) = typ
974            .keys
975            .iter()
976            .flatten()
977            .find(|key| **key >= typ.column_types.len())
978        {
979            return Err(TryFromProtoError::InvalidFieldError(format!(
980                "ProtoRelationDesc: key index {key} out of bounds for {} columns",
981                typ.column_types.len()
982            )));
983        }
984
985        // `metadata` Migration Logic: We wrote some `ProtoRelationDesc`s into Persist before the
986        // metadata field was added. If the field doesn't exist we fill it in with default values,
987        // and when converting into_proto we omit these fields so the serialized bytes roundtrip.
988        //
989        // Note: This logic needs to exist approximately forever.
990        let proto_metadata: Box<dyn Iterator<Item = _>> = if proto.metadata.is_empty() {
991            let val = ProtoColumnMetadata {
992                added: Some(RelationVersion::root().into_proto()),
993                dropped: None,
994            };
995            Box::new(itertools::repeat_n(val, proto.names.len()))
996        } else {
997            // Reject mismatched lengths explicitly rather than panicking via
998            // `zip_eq` below, since this branch is reachable from untrusted
999            // proto bytes.
1000            if proto.names.len() != proto.metadata.len() {
1001                return Err(TryFromProtoError::InvalidFieldError(format!(
1002                    "ProtoRelationDesc: names ({}) and metadata ({}) length mismatch",
1003                    proto.names.len(),
1004                    proto.metadata.len()
1005                )));
1006            }
1007            Box::new(proto.metadata.into_iter())
1008        };
1009
1010        let metadata = proto
1011            .names
1012            .into_iter()
1013            .zip_eq(proto_metadata)
1014            .enumerate()
1015            .map(|(idx, (name, metadata))| {
1016                let meta = ColumnMetadata {
1017                    name: name.into_rust()?,
1018                    typ_idx: idx,
1019                    added: metadata.added.into_rust_if_some("ColumnMetadata::added")?,
1020                    dropped: metadata.dropped.into_rust()?,
1021                };
1022                Ok::<_, TryFromProtoError>((ColumnIndex(idx), meta))
1023            })
1024            .collect::<Result<_, _>>()?;
1025
1026        Ok(RelationDesc { typ, metadata })
1027    }
1028}
1029
1030impl RelationDesc {
1031    /// Returns a [`RelationDescBuilder`] that can be used to construct a [`RelationDesc`].
1032    pub fn builder() -> RelationDescBuilder {
1033        RelationDescBuilder::default()
1034    }
1035
1036    /// Constructs a new `RelationDesc` that represents the empty relation
1037    /// with no columns and no keys.
1038    pub fn empty() -> Self {
1039        RelationDesc {
1040            typ: SqlRelationType::empty(),
1041            metadata: BTreeMap::default(),
1042        }
1043    }
1044
1045    /// Check if the `RelationDesc` is empty.
1046    pub fn is_empty(&self) -> bool {
1047        self == &Self::empty()
1048    }
1049
1050    /// Returns the number of columns in this [`RelationDesc`].
1051    pub fn len(&self) -> usize {
1052        self.typ().column_types.len()
1053    }
1054
1055    /// Constructs a new `RelationDesc` from a `SqlRelationType` and an iterator
1056    /// over column names.
1057    ///
1058    /// # Panics
1059    ///
1060    /// Panics if the arity of the `SqlRelationType` is not equal to the number of
1061    /// items in `names`.
1062    pub fn new<I, N>(typ: SqlRelationType, names: I) -> Self
1063    where
1064        I: IntoIterator<Item = N>,
1065        N: Into<ColumnName>,
1066    {
1067        let metadata: BTreeMap<_, _> = names
1068            .into_iter()
1069            .enumerate()
1070            .map(|(idx, name)| {
1071                let col_idx = ColumnIndex(idx);
1072                let metadata = ColumnMetadata {
1073                    name: name.into(),
1074                    typ_idx: idx,
1075                    added: RelationVersion::root(),
1076                    dropped: None,
1077                };
1078                (col_idx, metadata)
1079            })
1080            .collect();
1081
1082        // TODO(parkmycar): Add better validation here.
1083        assert_eq!(typ.column_types.len(), metadata.len());
1084
1085        RelationDesc { typ, metadata }
1086    }
1087
1088    pub fn from_names_and_types<I, T, N>(iter: I) -> Self
1089    where
1090        I: IntoIterator<Item = (N, T)>,
1091        T: Into<SqlColumnType>,
1092        N: Into<ColumnName>,
1093    {
1094        let (names, types): (Vec<_>, Vec<_>) = iter.into_iter().unzip();
1095        let types = types.into_iter().map(Into::into).collect();
1096        let typ = SqlRelationType::new(types);
1097        Self::new(typ, names)
1098    }
1099
1100    /// Concatenates a `RelationDesc` onto the end of this `RelationDesc`.
1101    ///
1102    /// # Panics
1103    ///
1104    /// Panics if either `self` or `other` have columns that were added at a
1105    /// [`RelationVersion`] other than [`RelationVersion::root`] or if any
1106    /// columns were dropped.
1107    ///
1108    /// TODO(parkmycar): Move this method to [`RelationDescBuilder`].
1109    pub fn concat(mut self, other: Self) -> Self {
1110        let self_len = self.typ.column_types.len();
1111
1112        for (typ, (_col_idx, meta)) in other.typ.column_types.into_iter().zip_eq(other.metadata) {
1113            assert_eq!(meta.added, RelationVersion::root());
1114            assert_none!(meta.dropped);
1115
1116            let new_idx = self.typ.columns().len();
1117            let new_meta = ColumnMetadata {
1118                name: meta.name,
1119                typ_idx: new_idx,
1120                added: RelationVersion::root(),
1121                dropped: None,
1122            };
1123
1124            self.typ.column_types.push(typ);
1125            let prev = self.metadata.insert(ColumnIndex(new_idx), new_meta);
1126
1127            assert_eq!(self.metadata.len(), self.typ.columns().len());
1128            assert_none!(prev);
1129        }
1130
1131        for k in other.typ.keys {
1132            let k = k.into_iter().map(|idx| idx + self_len).collect();
1133            self = self.with_key(k);
1134        }
1135        self
1136    }
1137
1138    /// Adds a new key for the relation.
1139    pub fn with_key(mut self, indices: Vec<usize>) -> Self {
1140        self.typ = self.typ.with_key(indices);
1141        self
1142    }
1143
1144    /// Drops all existing keys.
1145    pub fn without_keys(mut self) -> Self {
1146        self.typ.keys.clear();
1147        self
1148    }
1149
1150    /// Builds a new relation description with the column names replaced with
1151    /// new names.
1152    ///
1153    /// # Panics
1154    ///
1155    /// Panics if the arity of the relation type does not match the number of
1156    /// items in `names`.
1157    pub fn with_names<I, N>(self, names: I) -> Self
1158    where
1159        I: IntoIterator<Item = N>,
1160        N: Into<ColumnName>,
1161    {
1162        Self::new(self.typ, names)
1163    }
1164
1165    /// Computes the number of columns in the relation.
1166    pub fn arity(&self) -> usize {
1167        self.typ.arity()
1168    }
1169
1170    /// Returns the relation type underlying this relation description.
1171    pub fn typ(&self) -> &SqlRelationType {
1172        &self.typ
1173    }
1174
1175    /// Returns the owned relation type underlying this relation description.
1176    pub fn into_typ(self) -> SqlRelationType {
1177        self.typ
1178    }
1179
1180    /// Returns an iterator over the columns in this relation.
1181    pub fn iter(&self) -> impl Iterator<Item = (&ColumnName, &SqlColumnType)> {
1182        self.metadata.values().map(|meta| {
1183            let typ = &self.typ.columns()[meta.typ_idx];
1184            (&meta.name, typ)
1185        })
1186    }
1187
1188    /// Returns an iterator over the types of the columns in this relation.
1189    pub fn iter_types(&self) -> impl Iterator<Item = &SqlColumnType> {
1190        self.typ.column_types.iter()
1191    }
1192
1193    /// Returns an iterator over the names of the columns in this relation.
1194    pub fn iter_names(&self) -> impl Iterator<Item = &ColumnName> {
1195        self.metadata.values().map(|meta| &meta.name)
1196    }
1197
1198    /// Returns an iterator over the columns in this relation, with all their metadata.
1199    pub fn iter_all(&self) -> impl Iterator<Item = (&ColumnIndex, &ColumnName, &SqlColumnType)> {
1200        self.metadata.iter().map(|(col_idx, metadata)| {
1201            let col_typ = &self.typ.columns()[metadata.typ_idx];
1202            (col_idx, &metadata.name, col_typ)
1203        })
1204    }
1205
1206    /// Returns an iterator over the names of the columns in this relation that are "similar" to
1207    /// the provided `name`.
1208    pub fn iter_similar_names<'a>(
1209        &'a self,
1210        name: &'a ColumnName,
1211    ) -> impl Iterator<Item = &'a ColumnName> {
1212        self.iter_names().filter(|n| n.is_similar(name))
1213    }
1214
1215    /// Returns whether this [`RelationDesc`] contains a column at the specified index.
1216    pub fn contains_index(&self, idx: &ColumnIndex) -> bool {
1217        self.metadata.contains_key(idx)
1218    }
1219
1220    /// Finds a column by name.
1221    ///
1222    /// Returns the index and type of the column named `name`. If no column with
1223    /// the specified name exists, returns `None`. If multiple columns have the
1224    /// specified name, the leftmost column is returned.
1225    pub fn get_by_name(&self, name: &ColumnName) -> Option<(usize, &SqlColumnType)> {
1226        self.iter_names()
1227            .position(|n| n == name)
1228            .map(|i| (i, &self.typ.column_types[i]))
1229    }
1230
1231    /// Gets the name of the `i`th column.
1232    ///
1233    /// # Panics
1234    ///
1235    /// Panics if `i` is not a valid column index.
1236    ///
1237    /// TODO(parkmycar): Migrate all uses of this to [`RelationDesc::get_name_idx`].
1238    pub fn get_name(&self, i: usize) -> &ColumnName {
1239        // TODO(parkmycar): Refactor this to use `ColumnIndex`.
1240        self.get_name_idx(&ColumnIndex(i))
1241    }
1242
1243    /// Gets the name of the column at `idx`.
1244    ///
1245    /// # Panics
1246    ///
1247    /// Panics if no column exists at `idx`.
1248    pub fn get_name_idx(&self, idx: &ColumnIndex) -> &ColumnName {
1249        &self.metadata.get(idx).expect("should exist").name
1250    }
1251
1252    /// Mutably gets the name of the `i`th column.
1253    ///
1254    /// # Panics
1255    ///
1256    /// Panics if `i` is not a valid column index.
1257    pub fn get_name_mut(&mut self, i: usize) -> &mut ColumnName {
1258        // TODO(parkmycar): Refactor this to use `ColumnIndex`.
1259        &mut self
1260            .metadata
1261            .get_mut(&ColumnIndex(i))
1262            .expect("should exist")
1263            .name
1264    }
1265
1266    /// Gets the [`SqlColumnType`] of the column at `idx`.
1267    ///
1268    /// # Panics
1269    ///
1270    /// Panics if no column exists at `idx`.
1271    pub fn get_type(&self, idx: &ColumnIndex) -> &SqlColumnType {
1272        let typ_idx = self.metadata.get(idx).expect("should exist").typ_idx;
1273        &self.typ.column_types[typ_idx]
1274    }
1275
1276    /// Gets the name of the `i`th column if that column name is unambiguous.
1277    ///
1278    /// If at least one other column has the same name as the `i`th column,
1279    /// returns `None`. If the `i`th column has no name, returns `None`.
1280    ///
1281    /// # Panics
1282    ///
1283    /// Panics if `i` is not a valid column index.
1284    pub fn get_unambiguous_name(&self, i: usize) -> Option<&ColumnName> {
1285        let name = self.get_name(i);
1286        if self.iter_names().filter(|n| *n == name).count() == 1 {
1287            Some(name)
1288        } else {
1289            None
1290        }
1291    }
1292
1293    /// Verifies that `d` meets all of the constraints for the `i`th column of `self`.
1294    ///
1295    /// n.b. The only constraint MZ currently supports in NOT NULL, but this
1296    /// structure will be simple to extend.
1297    pub fn constraints_met(&self, i: usize, d: &Datum) -> Result<(), NotNullViolation> {
1298        let name = self.get_name(i);
1299        let typ = &self.typ.column_types[i];
1300        if d == &Datum::Null && !typ.nullable {
1301            Err(NotNullViolation(name.clone()))
1302        } else {
1303            Ok(())
1304        }
1305    }
1306
1307    /// Computes the differences between two [`RelationDesc`]s.
1308    ///
1309    /// Returns a rich diff describing which columns differ, and in what way.
1310    ///
1311    /// # Panics
1312    ///
1313    /// Panics if either `self` or `other` have columns that were added at a
1314    /// [`RelationVersion`] other than [`RelationVersion::root`] or if any
1315    /// columns were dropped.
1316    ///
1317    /// This simplifies things by allowing us to assume that `ColumnIndex`es are
1318    /// dense and that they match the indexes of `typ.columns()`. Without this
1319    /// we would, e.g., struggle comparing keys as those are in terms of
1320    /// `typ.columns()` indexes.
1321    pub fn diff(&self, other: &RelationDesc) -> RelationDescDiff {
1322        assert_eq!(self.metadata.len(), self.typ.columns().len());
1323        assert_eq!(other.metadata.len(), other.typ.columns().len());
1324        for (idx, meta) in self.metadata.iter().chain(other.metadata.iter()) {
1325            assert_eq!(meta.typ_idx, idx.0);
1326            assert_eq!(meta.added, RelationVersion::root());
1327            assert_none!(meta.dropped);
1328        }
1329
1330        let mut column_diffs = BTreeMap::new();
1331        let mut key_diff = None;
1332
1333        let left_arity = self.arity();
1334        let right_arity = other.arity();
1335        let common_arity = std::cmp::min(left_arity, right_arity);
1336
1337        for idx in 0..common_arity {
1338            let left_name = self.get_name(idx);
1339            let right_name = other.get_name(idx);
1340            let left_type = &self.typ.column_types[idx];
1341            let right_type = &other.typ.column_types[idx];
1342
1343            if left_name != right_name {
1344                let diff = ColumnDiff::NameMismatch {
1345                    left: left_name.clone(),
1346                    right: right_name.clone(),
1347                };
1348                column_diffs.insert(idx, diff);
1349            } else if left_type.scalar_type != right_type.scalar_type {
1350                let diff = ColumnDiff::TypeMismatch {
1351                    name: left_name.clone(),
1352                    left: left_type.scalar_type.clone(),
1353                    right: right_type.scalar_type.clone(),
1354                };
1355                column_diffs.insert(idx, diff);
1356            } else if left_type.nullable != right_type.nullable {
1357                let diff = ColumnDiff::NullabilityMismatch {
1358                    name: left_name.clone(),
1359                    left: left_type.nullable,
1360                    right: right_type.nullable,
1361                };
1362                column_diffs.insert(idx, diff);
1363            }
1364        }
1365
1366        for idx in common_arity..left_arity {
1367            let diff = ColumnDiff::Missing {
1368                name: self.get_name(idx).clone(),
1369            };
1370            column_diffs.insert(idx, diff);
1371        }
1372
1373        for idx in common_arity..right_arity {
1374            let diff = ColumnDiff::Extra {
1375                name: other.get_name(idx).clone(),
1376            };
1377            column_diffs.insert(idx, diff);
1378        }
1379
1380        let left_keys: BTreeSet<_> = self.typ.keys.iter().collect();
1381        let right_keys: BTreeSet<_> = other.typ.keys.iter().collect();
1382        if left_keys != right_keys {
1383            let column_names = |desc: &RelationDesc, keys: BTreeSet<&Vec<usize>>| {
1384                keys.iter()
1385                    .map(|key| key.iter().map(|&idx| desc.get_name(idx).clone()).collect())
1386                    .collect()
1387            };
1388            key_diff = Some(KeyDiff {
1389                left: column_names(self, left_keys),
1390                right: column_names(other, right_keys),
1391            });
1392        }
1393
1394        RelationDescDiff {
1395            column_diffs,
1396            key_diff,
1397        }
1398    }
1399
1400    /// Creates a new [`RelationDesc`] retaining only the columns specified in `demands`.
1401    pub fn apply_demand(&self, demands: &BTreeSet<usize>) -> RelationDesc {
1402        // This filters `metadata` by raw ColumnIndex but `typ` by position,
1403        // which only agree when the desc is dense. Every desc constructible
1404        // today is (schema history is add-only), but a dropped column would
1405        // desync the two and silently attach types, statistics, and filter
1406        // specs to the wrong columns downstream.
1407        debug_assert!(
1408            self.metadata
1409                .iter()
1410                .enumerate()
1411                .all(|(pos, (idx, meta))| idx.0 == pos && meta.typ_idx == pos),
1412            "apply_demand requires a dense RelationDesc (ColumnIndex == typ_idx): {:?}",
1413            self.metadata,
1414        );
1415        let mut new_desc = self.clone();
1416
1417        // Update ColumnMetadata.
1418        let mut removed = 0;
1419        new_desc.metadata.retain(|idx, metadata| {
1420            let retain = demands.contains(&idx.0);
1421            if !retain {
1422                removed += 1;
1423            } else {
1424                metadata.typ_idx -= removed;
1425            }
1426            retain
1427        });
1428
1429        // Update SqlColumnType.
1430        let mut idx = 0;
1431        new_desc.typ.column_types.retain(|_| {
1432            let keep = demands.contains(&idx);
1433            idx += 1;
1434            keep
1435        });
1436
1437        new_desc
1438    }
1439}
1440
1441#[cfg(any(test, feature = "proptest"))]
1442impl Arbitrary for RelationDesc {
1443    type Parameters = ();
1444    type Strategy = BoxedStrategy<RelationDesc>;
1445
1446    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
1447        let mut weights = vec![(100, Just(0..4)), (50, Just(4..8)), (25, Just(8..16))];
1448        if std::env::var("PROPTEST_LARGE_DATA").is_ok() {
1449            weights.extend([
1450                (12, Just(16..32)),
1451                (6, Just(32..64)),
1452                (3, Just(64..128)),
1453                (1, Just(128..256)),
1454            ]);
1455        }
1456        let num_columns = Union::new_weighted(weights);
1457
1458        num_columns.prop_flat_map(arb_relation_desc).boxed()
1459    }
1460}
1461
1462/// Returns a [`Strategy`] that generates an arbitrary [`RelationDesc`] with a number columns
1463/// within the range provided.
1464#[cfg(any(test, feature = "proptest"))]
1465pub fn arb_relation_desc(num_cols: std::ops::Range<usize>) -> impl Strategy<Value = RelationDesc> {
1466    proptest::collection::btree_map(any::<ColumnName>(), any::<SqlColumnType>(), num_cols)
1467        .prop_map(RelationDesc::from_names_and_types)
1468}
1469
1470/// Returns a [`Strategy`] that generates a projection of the provided [`RelationDesc`].
1471#[cfg(any(test, feature = "proptest"))]
1472pub fn arb_relation_desc_projection(desc: RelationDesc) -> impl Strategy<Value = RelationDesc> {
1473    let mask: Vec<_> = (0..desc.len()).map(|_| any::<bool>()).collect();
1474    mask.prop_map(move |mask| {
1475        let demands: BTreeSet<_> = mask
1476            .into_iter()
1477            .enumerate()
1478            .filter_map(|(idx, keep)| keep.then_some(idx))
1479            .collect();
1480        desc.apply_demand(&demands)
1481    })
1482}
1483
1484impl IntoIterator for RelationDesc {
1485    type Item = (ColumnName, SqlColumnType);
1486    type IntoIter = Box<dyn Iterator<Item = (ColumnName, SqlColumnType)>>;
1487
1488    fn into_iter(self) -> Self::IntoIter {
1489        let iter = self
1490            .metadata
1491            .into_values()
1492            .zip_eq(self.typ.column_types)
1493            .map(|(meta, typ)| (meta.name, typ));
1494        Box::new(iter)
1495    }
1496}
1497
1498/// Returns a [`Strategy`] that yields arbitrary [`Row`]s for the provided [`RelationDesc`].
1499#[cfg(any(test, feature = "proptest"))]
1500pub fn arb_row_for_relation(desc: &RelationDesc) -> impl Strategy<Value = Row> + use<> {
1501    let datums: Vec<_> = desc
1502        .typ()
1503        .columns()
1504        .iter()
1505        .cloned()
1506        .map(arb_datum_for_column)
1507        .collect();
1508    datums.prop_map(|x| Row::pack(x.iter().map(Datum::from)))
1509}
1510
1511/// Expression violated not-null constraint on named column
1512#[derive(Debug, PartialEq, Eq)]
1513pub struct NotNullViolation(pub ColumnName);
1514
1515impl fmt::Display for NotNullViolation {
1516    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1517        write!(
1518            f,
1519            "null value in column {} violates not-null constraint",
1520            self.0.quoted()
1521        )
1522    }
1523}
1524
1525/// The result of comparing two [`RelationDesc`]s.
1526#[derive(Debug, Clone, PartialEq, Eq)]
1527pub struct RelationDescDiff {
1528    /// Column differences, keyed by column index.
1529    pub column_diffs: BTreeMap<usize, ColumnDiff>,
1530    /// Key differences, if any.
1531    pub key_diff: Option<KeyDiff>,
1532}
1533
1534impl RelationDescDiff {
1535    /// Returns whether the diff contains any differences.
1536    pub fn is_empty(&self) -> bool {
1537        self.column_diffs.is_empty() && self.key_diff.is_none()
1538    }
1539}
1540
1541/// A difference in a column between two [`RelationDesc`]s.
1542#[derive(Debug, Clone, PartialEq, Eq)]
1543pub enum ColumnDiff {
1544    /// Column exists only in the left relation.
1545    Missing { name: ColumnName },
1546    /// Column exists only in the right relation.
1547    Extra { name: ColumnName },
1548    /// Columns have different types.
1549    TypeMismatch {
1550        name: ColumnName,
1551        left: SqlScalarType,
1552        right: SqlScalarType,
1553    },
1554    /// Columns have different nullability.
1555    NullabilityMismatch {
1556        name: ColumnName,
1557        left: bool,
1558        right: bool,
1559    },
1560    /// Columns have different names.
1561    NameMismatch { left: ColumnName, right: ColumnName },
1562}
1563
1564/// A difference in the keys of two [`RelationDesc`]s.
1565#[derive(Debug, Clone, PartialEq, Eq)]
1566pub struct KeyDiff {
1567    /// Keys of the left relation.
1568    pub left: BTreeSet<Vec<ColumnName>>,
1569    /// Keys of the right relation.
1570    pub right: BTreeSet<Vec<ColumnName>>,
1571}
1572
1573/// A builder for a [`RelationDesc`].
1574#[derive(Clone, Default, Debug, PartialEq, Eq)]
1575pub struct RelationDescBuilder {
1576    /// Columns of the relation.
1577    columns: Vec<(ColumnName, SqlColumnType)>,
1578    /// Sets of indices that are "keys" for the collection.
1579    keys: Vec<Vec<usize>>,
1580}
1581
1582impl RelationDescBuilder {
1583    /// Appends a column with the specified name and type.
1584    pub fn with_column<N: Into<ColumnName>>(
1585        mut self,
1586        name: N,
1587        ty: SqlColumnType,
1588    ) -> RelationDescBuilder {
1589        let name = name.into();
1590        self.columns.push((name, ty));
1591        self
1592    }
1593
1594    /// Appends the provided columns to the builder.
1595    pub fn with_columns<I, T, N>(mut self, iter: I) -> Self
1596    where
1597        I: IntoIterator<Item = (N, T)>,
1598        T: Into<SqlColumnType>,
1599        N: Into<ColumnName>,
1600    {
1601        self.columns
1602            .extend(iter.into_iter().map(|(name, ty)| (name.into(), ty.into())));
1603        self
1604    }
1605
1606    /// Adds a new key for the relation.
1607    pub fn with_key(mut self, mut indices: Vec<usize>) -> RelationDescBuilder {
1608        indices.sort_unstable();
1609        if !self.keys.contains(&indices) {
1610            self.keys.push(indices);
1611        }
1612        self
1613    }
1614
1615    /// Removes all previously inserted keys.
1616    pub fn without_keys(mut self) -> RelationDescBuilder {
1617        self.keys.clear();
1618        assert_eq!(self.keys.len(), 0);
1619        self
1620    }
1621
1622    /// Concatenates a [`RelationDescBuilder`] onto the end of this [`RelationDescBuilder`].
1623    pub fn concat(mut self, other: Self) -> Self {
1624        let self_len = self.columns.len();
1625
1626        self.columns.extend(other.columns);
1627        for k in other.keys {
1628            let k = k.into_iter().map(|idx| idx + self_len).collect();
1629            self = self.with_key(k);
1630        }
1631
1632        self
1633    }
1634
1635    /// Finish the builder, returning a [`RelationDesc`].
1636    pub fn finish(self) -> RelationDesc {
1637        let mut desc = RelationDesc::from_names_and_types(self.columns);
1638        desc.typ = desc.typ.with_keys(self.keys);
1639        desc
1640    }
1641}
1642
1643/// Describes a [`RelationDesc`] at a specific version of a [`VersionedRelationDesc`].
1644#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize)]
1645pub enum RelationVersionSelector {
1646    Specific(RelationVersion),
1647    Latest,
1648}
1649
1650impl RelationVersionSelector {
1651    pub fn specific(version: u64) -> Self {
1652        RelationVersionSelector::Specific(RelationVersion(version))
1653    }
1654}
1655
1656/// A wrapper around [`RelationDesc`] that provides an interface for adding
1657/// columns and generating new versions.
1658///
1659/// TODO(parkmycar): Using an immutable data structure for RelationDesc would
1660/// be great.
1661#[derive(Debug, Clone, Serialize)]
1662pub struct VersionedRelationDesc {
1663    inner: RelationDesc,
1664}
1665
1666impl VersionedRelationDesc {
1667    pub fn new(inner: RelationDesc) -> Self {
1668        VersionedRelationDesc { inner }
1669    }
1670
1671    /// Adds a new column to this [`RelationDesc`], creating a new version of the [`RelationDesc`].
1672    ///
1673    /// # Panics
1674    ///
1675    /// * Panics if a column with `name` already exists that hasn't been dropped.
1676    ///
1677    /// Note: For building a [`RelationDesc`] see [`RelationDescBuilder::with_column`].
1678    #[must_use]
1679    pub fn add_column<N, T>(&mut self, name: N, typ: T) -> RelationVersion
1680    where
1681        N: Into<ColumnName>,
1682        T: Into<SqlColumnType>,
1683    {
1684        let latest_version = self.latest_version();
1685        let new_version = latest_version.bump();
1686
1687        let name = name.into();
1688        let existing = self
1689            .inner
1690            .metadata
1691            .iter()
1692            .find(|(_, meta)| meta.name == name && meta.dropped.is_none());
1693        if let Some(existing) = existing {
1694            panic!("column named '{name}' already exists! {existing:?}");
1695        }
1696
1697        let next_idx = self.inner.metadata.len();
1698        let col_meta = ColumnMetadata {
1699            name,
1700            typ_idx: next_idx,
1701            added: new_version,
1702            dropped: None,
1703        };
1704
1705        self.inner.typ.column_types.push(typ.into());
1706        let prev = self.inner.metadata.insert(ColumnIndex(next_idx), col_meta);
1707
1708        assert_none!(prev, "column index overlap!");
1709        self.validate();
1710
1711        new_version
1712    }
1713
1714    /// Drops the column `name` from this [`RelationDesc`]. If there are multiple columns with
1715    /// `name` drops the left-most one that hasn't already been dropped.
1716    ///
1717    /// TODO(parkmycar): Add handling for dropping a column that is currently used as a key.
1718    ///
1719    /// # Panics
1720    ///
1721    /// Panics if a column with `name` does not exist or the dropped column was used as a key.
1722    #[must_use]
1723    pub fn drop_column<N>(&mut self, name: N) -> RelationVersion
1724    where
1725        N: Into<ColumnName>,
1726    {
1727        let name = name.into();
1728        let latest_version = self.latest_version();
1729        let new_version = latest_version.bump();
1730
1731        let col = self
1732            .inner
1733            .metadata
1734            .values_mut()
1735            .find(|meta| meta.name == name && meta.dropped.is_none())
1736            .expect("column to exist");
1737
1738        // Make sure the column hadn't been previously dropped.
1739        assert_none!(col.dropped, "column was already dropped");
1740        col.dropped = Some(new_version);
1741
1742        // Make sure the column isn't being used as a key.
1743        let dropped_key = self
1744            .inner
1745            .typ
1746            .keys
1747            .iter()
1748            .any(|keys| keys.contains(&col.typ_idx));
1749        assert!(!dropped_key, "column being dropped was used as a key");
1750
1751        self.validate();
1752        new_version
1753    }
1754
1755    /// Returns the [`RelationDesc`] at the latest version.
1756    pub fn latest(&self) -> RelationDesc {
1757        self.inner.clone()
1758    }
1759
1760    /// Returns this [`RelationDesc`] at the specified version.
1761    pub fn at_version(&self, version: RelationVersionSelector) -> RelationDesc {
1762        // Get all of the changes from the start, up to whatever version was requested.
1763        let up_to_version = match version {
1764            RelationVersionSelector::Latest => RelationVersion(u64::MAX),
1765            RelationVersionSelector::Specific(v) => v,
1766        };
1767
1768        let valid_columns = self.inner.metadata.iter().filter(|(_col_idx, meta)| {
1769            let added = meta.added <= up_to_version;
1770            let dropped = meta
1771                .dropped
1772                .map(|dropped_at| up_to_version >= dropped_at)
1773                .unwrap_or(false);
1774
1775            added && !dropped
1776        });
1777
1778        let mut column_types = Vec::new();
1779        let mut column_metas = BTreeMap::new();
1780
1781        // N.B. At this point we need to be careful because col_idx might not
1782        // equal typ_idx.
1783        //
1784        // For example, consider columns "a", "b", and "c" with indexes 0, 1,
1785        // and 2. If we drop column "b" then we'll have "a" and "c" with column
1786        // indexes 0 and 2, but their indices in SqlRelationType will be 0 and 1.
1787        for (col_idx, meta) in valid_columns {
1788            let new_meta = ColumnMetadata {
1789                name: meta.name.clone(),
1790                typ_idx: column_types.len(),
1791                added: meta.added.clone(),
1792                dropped: meta.dropped.clone(),
1793            };
1794            column_types.push(self.inner.typ.columns()[meta.typ_idx].clone());
1795            column_metas.insert(*col_idx, new_meta);
1796        }
1797
1798        // Remap keys in case a column with an index less than that of a key was
1799        // dropped.
1800        //
1801        // For example, consider columns "a", "b", and "c" where "a" and "c" are
1802        // keys and "b" was dropped.
1803        let keys = self
1804            .inner
1805            .typ
1806            .keys
1807            .iter()
1808            .map(|keys| {
1809                keys.iter()
1810                    .map(|key_idx| {
1811                        let metadata = column_metas
1812                            .get(&ColumnIndex(*key_idx))
1813                            .expect("found key for column that doesn't exist");
1814                        metadata.typ_idx
1815                    })
1816                    .collect()
1817            })
1818            .collect();
1819
1820        let relation_type = SqlRelationType { column_types, keys };
1821
1822        RelationDesc {
1823            typ: relation_type,
1824            metadata: column_metas,
1825        }
1826    }
1827
1828    pub fn latest_version(&self) -> RelationVersion {
1829        self.inner
1830            .metadata
1831            .values()
1832            // N.B. Dropped is always greater than added.
1833            .map(|meta| meta.dropped.unwrap_or(meta.added))
1834            .max()
1835            // If there aren't any columns we're implicitly the root version.
1836            .unwrap_or_else(RelationVersion::root)
1837    }
1838
1839    /// Validates internal contraints of the [`RelationDesc`] are correct.
1840    ///
1841    /// # Panics
1842    ///
1843    /// Panics if a constraint is not satisfied.
1844    fn validate(&self) {
1845        fn validate_inner(desc: &RelationDesc) -> Result<(), anyhow::Error> {
1846            if desc.typ.column_types.len() != desc.metadata.len() {
1847                anyhow::bail!("mismatch between number of types and metadatas");
1848            }
1849
1850            for (col_idx, meta) in &desc.metadata {
1851                if col_idx.0 > desc.metadata.len() {
1852                    anyhow::bail!("column index out of bounds");
1853                }
1854                if meta.added >= meta.dropped.unwrap_or(RelationVersion(u64::MAX)) {
1855                    anyhow::bail!("column was added after it was dropped?");
1856                }
1857                if desc.typ().columns().get(meta.typ_idx).is_none() {
1858                    anyhow::bail!("typ_idx incorrect");
1859                }
1860            }
1861
1862            for keys in &desc.typ.keys {
1863                for key in keys {
1864                    if *key >= desc.typ.column_types.len() {
1865                        anyhow::bail!("key index was out of bounds!");
1866                    }
1867                }
1868            }
1869
1870            // Every version past the root is stamped exactly once, by the add or the
1871            // drop that created it, so a column that was added and later dropped
1872            // accounts for two of them and both have to be counted. Collapsing to
1873            // `dropped.unwrap_or(added)` would lose the add's version, and with it any
1874            // desc where a non-root column was later dropped.
1875            let versions = desc
1876                .metadata
1877                .values()
1878                .flat_map(|meta| [Some(meta.added), meta.dropped])
1879                .flatten()
1880                // The root version is the one version many columns can share.
1881                .filter(|version| *version != RelationVersion::root());
1882            let mut max = 0;
1883            let mut sum = 0;
1884            for version in versions {
1885                max = std::cmp::max(max, version.0);
1886                sum += version.0;
1887            }
1888
1889            // Other than RelationVersion(0), we should never have duplicate
1890            // versions and they should always increase by 1. In other words, the
1891            // sum of all RelationVersions should be the sum of [0, max].
1892            //
1893            // N.B. n * (n + 1) / 2 = sum of [0, n]
1894            //
1895            // While I normally don't like tricks like this, it allows us to
1896            // validate that our column versions are correct in O(n) time and
1897            // without allocations.
1898            if sum != (max * (max + 1) / 2) {
1899                anyhow::bail!("there is a duplicate or missing relation version");
1900            }
1901
1902            Ok(())
1903        }
1904
1905        assert_ok!(validate_inner(&self.inner), "validate failed! {self:?}");
1906    }
1907}
1908
1909/// Diffs that can be generated proptest and applied to a [`RelationDesc`] to
1910/// exercise schema migrations.
1911#[derive(Debug)]
1912#[cfg(any(test, feature = "proptest"))]
1913pub enum PropRelationDescDiff {
1914    AddColumn {
1915        name: ColumnName,
1916        typ: SqlColumnType,
1917    },
1918    DropColumn {
1919        name: ColumnName,
1920    },
1921    ToggleNullability {
1922        name: ColumnName,
1923    },
1924    ChangeType {
1925        name: ColumnName,
1926        typ: SqlColumnType,
1927    },
1928}
1929
1930#[cfg(any(test, feature = "proptest"))]
1931impl PropRelationDescDiff {
1932    pub fn apply(self, desc: &mut RelationDesc) {
1933        match self {
1934            PropRelationDescDiff::AddColumn { name, typ } => {
1935                let new_idx = desc.metadata.len();
1936                let meta = ColumnMetadata {
1937                    name,
1938                    typ_idx: new_idx,
1939                    added: RelationVersion(0),
1940                    dropped: None,
1941                };
1942                let prev = desc.metadata.insert(ColumnIndex(new_idx), meta);
1943                desc.typ.column_types.push(typ);
1944
1945                assert_none!(prev);
1946                assert_eq!(desc.metadata.len(), desc.typ.column_types.len());
1947            }
1948            PropRelationDescDiff::DropColumn { name } => {
1949                let next_version = desc
1950                    .metadata
1951                    .values()
1952                    .map(|meta| meta.dropped.unwrap_or(meta.added))
1953                    .max()
1954                    .unwrap_or_else(RelationVersion::root)
1955                    .bump();
1956                let Some(metadata) = desc.metadata.values_mut().find(|meta| meta.name == name)
1957                else {
1958                    return;
1959                };
1960                if metadata.dropped.is_none() {
1961                    metadata.dropped = Some(next_version);
1962                }
1963            }
1964            PropRelationDescDiff::ToggleNullability { name } => {
1965                let Some((pos, _)) = desc.get_by_name(&name) else {
1966                    return;
1967                };
1968                let col_type = desc
1969                    .typ
1970                    .column_types
1971                    .get_mut(pos)
1972                    .expect("ColumnNames and SqlColumnTypes out of sync!");
1973                col_type.nullable = !col_type.nullable;
1974            }
1975            PropRelationDescDiff::ChangeType { name, typ } => {
1976                let Some((pos, _)) = desc.get_by_name(&name) else {
1977                    return;
1978                };
1979                let col_type = desc
1980                    .typ
1981                    .column_types
1982                    .get_mut(pos)
1983                    .expect("ColumnNames and SqlColumnTypes out of sync!");
1984                *col_type = typ;
1985            }
1986        }
1987    }
1988}
1989
1990/// Generates a set of [`PropRelationDescDiff`]s based on some source [`RelationDesc`].
1991#[cfg(any(test, feature = "proptest"))]
1992pub fn arb_relation_desc_diff(
1993    source: &RelationDesc,
1994) -> impl Strategy<Value = Vec<PropRelationDescDiff>> + use<> {
1995    let source = Rc::new(source.clone());
1996    let num_source_columns = source.typ.columns().len();
1997
1998    let num_add_columns = Union::new_weighted(vec![(100, Just(0..8)), (1, Just(8..64))]);
1999    let add_columns_strat = num_add_columns
2000        .prop_flat_map(|num_columns| {
2001            proptest::collection::vec((any::<ColumnName>(), any::<SqlColumnType>()), num_columns)
2002        })
2003        .prop_map(|cols| {
2004            cols.into_iter()
2005                .map(|(name, typ)| PropRelationDescDiff::AddColumn { name, typ })
2006                .collect::<Vec<_>>()
2007        });
2008
2009    // If the source RelationDesc is empty there is nothing else to do.
2010    if num_source_columns == 0 {
2011        return add_columns_strat.boxed();
2012    }
2013
2014    let source_ = Rc::clone(&source);
2015    let drop_columns_strat = (0..num_source_columns).prop_perturb(move |num_columns, mut rng| {
2016        let mut set = BTreeSet::default();
2017        for _ in 0..num_columns {
2018            let col_idx = rng.random_range(0..num_source_columns);
2019            set.insert(source_.get_name(col_idx).clone());
2020        }
2021        set.into_iter()
2022            .map(|name| PropRelationDescDiff::DropColumn { name })
2023            .collect::<Vec<_>>()
2024    });
2025
2026    let source_ = Rc::clone(&source);
2027    let toggle_nullability_strat =
2028        (0..num_source_columns).prop_perturb(move |num_columns, mut rng| {
2029            let mut set = BTreeSet::default();
2030            for _ in 0..num_columns {
2031                let col_idx = rng.random_range(0..num_source_columns);
2032                set.insert(source_.get_name(col_idx).clone());
2033            }
2034            set.into_iter()
2035                .map(|name| PropRelationDescDiff::ToggleNullability { name })
2036                .collect::<Vec<_>>()
2037        });
2038
2039    let source_ = Rc::clone(&source);
2040    let change_type_strat = (0..num_source_columns)
2041        .prop_perturb(move |num_columns, mut rng| {
2042            let mut set = BTreeSet::default();
2043            for _ in 0..num_columns {
2044                let col_idx = rng.random_range(0..num_source_columns);
2045                set.insert(source_.get_name(col_idx).clone());
2046            }
2047            set
2048        })
2049        .prop_flat_map(|cols| {
2050            proptest::collection::vec(any::<SqlColumnType>(), cols.len())
2051                .prop_map(move |types| (cols.clone(), types))
2052        })
2053        .prop_map(|(cols, types)| {
2054            cols.into_iter()
2055                .zip_eq(types)
2056                .map(|(name, typ)| PropRelationDescDiff::ChangeType { name, typ })
2057                .collect::<Vec<_>>()
2058        });
2059
2060    (
2061        add_columns_strat,
2062        drop_columns_strat,
2063        toggle_nullability_strat,
2064        change_type_strat,
2065    )
2066        .prop_map(|(adds, drops, toggles, changes)| {
2067            adds.into_iter()
2068                .chain(drops)
2069                .chain(toggles)
2070                .chain(changes)
2071                .collect::<Vec<_>>()
2072        })
2073        .prop_shuffle()
2074        .boxed()
2075}
2076
2077#[cfg(test)]
2078mod tests {
2079    use super::*;
2080    use prost::Message;
2081
2082    /// `apply_demand`, and the stats and filter-spec plumbing downstream of
2083    /// it, require dense descs. A desc with a dropped column must trip the
2084    /// assertion rather than silently misattach columns.
2085    #[mz_ore::test]
2086    #[should_panic(expected = "dense RelationDesc")]
2087    fn apply_demand_rejects_non_dense_desc() {
2088        let desc = RelationDesc::builder()
2089            .with_column("a", SqlScalarType::Int32.nullable(false))
2090            .with_column("b", SqlScalarType::Int32.nullable(false))
2091            .with_column("c", SqlScalarType::Int32.nullable(false))
2092            .finish();
2093        let mut versioned = VersionedRelationDesc::new(desc);
2094        let version = versioned.drop_column("b");
2095        let desc = versioned.at_version(RelationVersionSelector::Specific(version));
2096        let _ = desc.apply_demand(&BTreeSet::from([0]));
2097    }
2098
2099    #[mz_ore::test]
2100    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `pipe2` on OS `linux`
2101    fn smoktest_at_version() {
2102        let desc = RelationDesc::builder()
2103            .with_column("a", SqlScalarType::Bool.nullable(true))
2104            .with_column("z", SqlScalarType::String.nullable(false))
2105            .finish();
2106
2107        let mut versioned_desc = VersionedRelationDesc {
2108            inner: desc.clone(),
2109        };
2110        versioned_desc.validate();
2111
2112        let latest = versioned_desc.at_version(RelationVersionSelector::Latest);
2113        assert_eq!(desc, latest);
2114
2115        let v0 = versioned_desc.at_version(RelationVersionSelector::specific(0));
2116        assert_eq!(desc, v0);
2117
2118        let v3 = versioned_desc.at_version(RelationVersionSelector::specific(3));
2119        assert_eq!(desc, v3);
2120
2121        let v1 = versioned_desc.add_column("b", SqlScalarType::Bytes.nullable(false));
2122        assert_eq!(v1, RelationVersion(1));
2123
2124        let v1 = versioned_desc.at_version(RelationVersionSelector::Specific(v1));
2125        insta::assert_json_snapshot!(v1.metadata, @r###"
2126        {
2127          "0": {
2128            "name": "a",
2129            "typ_idx": 0,
2130            "added": 0,
2131            "dropped": null
2132          },
2133          "1": {
2134            "name": "z",
2135            "typ_idx": 1,
2136            "added": 0,
2137            "dropped": null
2138          },
2139          "2": {
2140            "name": "b",
2141            "typ_idx": 2,
2142            "added": 1,
2143            "dropped": null
2144          }
2145        }
2146        "###);
2147
2148        // Check that V0 doesn't show the new column.
2149        let v0_b = versioned_desc.at_version(RelationVersionSelector::specific(0));
2150        assert!(v0.iter().eq(v0_b.iter()));
2151
2152        let v2 = versioned_desc.drop_column("z");
2153        assert_eq!(v2, RelationVersion(2));
2154
2155        let v2 = versioned_desc.at_version(RelationVersionSelector::Specific(v2));
2156        insta::assert_json_snapshot!(v2.metadata, @r###"
2157        {
2158          "0": {
2159            "name": "a",
2160            "typ_idx": 0,
2161            "added": 0,
2162            "dropped": null
2163          },
2164          "2": {
2165            "name": "b",
2166            "typ_idx": 1,
2167            "added": 1,
2168            "dropped": null
2169          }
2170        }
2171        "###);
2172
2173        // Check that V0 and V1 are still correct.
2174        let v0_c = versioned_desc.at_version(RelationVersionSelector::specific(0));
2175        assert!(v0.iter().eq(v0_c.iter()));
2176
2177        let v1_b = versioned_desc.at_version(RelationVersionSelector::specific(1));
2178        assert!(v1.iter().eq(v1_b.iter()));
2179
2180        insta::assert_json_snapshot!(versioned_desc.inner.metadata, @r###"
2181        {
2182          "0": {
2183            "name": "a",
2184            "typ_idx": 0,
2185            "added": 0,
2186            "dropped": null
2187          },
2188          "1": {
2189            "name": "z",
2190            "typ_idx": 1,
2191            "added": 0,
2192            "dropped": 2
2193          },
2194          "2": {
2195            "name": "b",
2196            "typ_idx": 2,
2197            "added": 1,
2198            "dropped": null
2199          }
2200        }
2201        "###);
2202    }
2203
2204    #[mz_ore::test]
2205    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `pipe2` on OS `linux`
2206    fn test_dropping_columns_with_keys() {
2207        let desc = RelationDesc::builder()
2208            .with_column("a", SqlScalarType::Bool.nullable(true))
2209            .with_column("z", SqlScalarType::String.nullable(false))
2210            .with_key(vec![1])
2211            .finish();
2212
2213        let mut versioned_desc = VersionedRelationDesc {
2214            inner: desc.clone(),
2215        };
2216        versioned_desc.validate();
2217
2218        let v1 = versioned_desc.drop_column("a");
2219        assert_eq!(v1, RelationVersion(1));
2220
2221        // Make sure the key index for 'z' got remapped since 'a' was dropped.
2222        let v1 = versioned_desc.at_version(RelationVersionSelector::Specific(v1));
2223        insta::assert_json_snapshot!(v1, @r###"
2224        {
2225          "typ": {
2226            "column_types": [
2227              {
2228                "scalar_type": "String",
2229                "nullable": false
2230              }
2231            ],
2232            "keys": [
2233              [
2234                0
2235              ]
2236            ]
2237          },
2238          "metadata": {
2239            "1": {
2240              "name": "z",
2241              "typ_idx": 0,
2242              "added": 0,
2243              "dropped": null
2244            }
2245          }
2246        }
2247        "###);
2248
2249        // Make sure the key index of 'z' is correct when all columns are present.
2250        let v0 = versioned_desc.at_version(RelationVersionSelector::specific(0));
2251        insta::assert_json_snapshot!(v0, @r###"
2252        {
2253          "typ": {
2254            "column_types": [
2255              {
2256                "scalar_type": "Bool",
2257                "nullable": true
2258              },
2259              {
2260                "scalar_type": "String",
2261                "nullable": false
2262              }
2263            ],
2264            "keys": [
2265              [
2266                1
2267              ]
2268            ]
2269          },
2270          "metadata": {
2271            "0": {
2272              "name": "a",
2273              "typ_idx": 0,
2274              "added": 0,
2275              "dropped": 1
2276            },
2277            "1": {
2278              "name": "z",
2279              "typ_idx": 1,
2280              "added": 0,
2281              "dropped": null
2282            }
2283          }
2284        }
2285        "###);
2286    }
2287
2288    #[mz_ore::test]
2289    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `pipe2` on OS `linux`
2290    fn roundtrip_relation_desc_without_metadata() {
2291        let typ = ProtoRelationType {
2292            column_types: vec![
2293                SqlScalarType::String.nullable(false).into_proto(),
2294                SqlScalarType::Bool.nullable(true).into_proto(),
2295            ],
2296            keys: vec![],
2297        };
2298        let proto = ProtoRelationDesc {
2299            typ: Some(typ),
2300            names: vec![
2301                ColumnName("a".into()).into_proto(),
2302                ColumnName("b".into()).into_proto(),
2303            ],
2304            metadata: vec![],
2305        };
2306        let desc: RelationDesc = proto.into_rust().unwrap();
2307
2308        insta::assert_json_snapshot!(desc, @r###"
2309        {
2310          "typ": {
2311            "column_types": [
2312              {
2313                "scalar_type": "String",
2314                "nullable": false
2315              },
2316              {
2317                "scalar_type": "Bool",
2318                "nullable": true
2319              }
2320            ],
2321            "keys": []
2322          },
2323          "metadata": {
2324            "0": {
2325              "name": "a",
2326              "typ_idx": 0,
2327              "added": 0,
2328              "dropped": null
2329            },
2330            "1": {
2331              "name": "b",
2332              "typ_idx": 1,
2333              "added": 0,
2334              "dropped": null
2335            }
2336          }
2337        }
2338        "###);
2339    }
2340
2341    #[mz_ore::test]
2342    #[should_panic(expected = "column named 'a' already exists!")]
2343    fn test_add_column_with_same_name_panics() {
2344        let desc = RelationDesc::builder()
2345            .with_column("a", SqlScalarType::Bool.nullable(true))
2346            .finish();
2347        let mut versioned = VersionedRelationDesc::new(desc);
2348
2349        let _ = versioned.add_column("a", SqlScalarType::String.nullable(false));
2350    }
2351
2352    #[mz_ore::test]
2353    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `pipe2` on OS `linux`
2354    fn test_add_column_with_same_name_prev_dropped() {
2355        let desc = RelationDesc::builder()
2356            .with_column("a", SqlScalarType::Bool.nullable(true))
2357            .finish();
2358        let mut versioned = VersionedRelationDesc::new(desc);
2359
2360        let v1 = versioned.drop_column("a");
2361        let v1 = versioned.at_version(RelationVersionSelector::Specific(v1));
2362        insta::assert_json_snapshot!(v1, @r###"
2363        {
2364          "typ": {
2365            "column_types": [],
2366            "keys": []
2367          },
2368          "metadata": {}
2369        }
2370        "###);
2371
2372        let v2 = versioned.add_column("a", SqlScalarType::String.nullable(false));
2373        let v2 = versioned.at_version(RelationVersionSelector::Specific(v2));
2374        insta::assert_json_snapshot!(v2, @r###"
2375        {
2376          "typ": {
2377            "column_types": [
2378              {
2379                "scalar_type": "String",
2380                "nullable": false
2381              }
2382            ],
2383            "keys": []
2384          },
2385          "metadata": {
2386            "1": {
2387              "name": "a",
2388              "typ_idx": 0,
2389              "added": 2,
2390              "dropped": null
2391            }
2392          }
2393        }
2394        "###);
2395    }
2396
2397    #[mz_ore::test]
2398    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `pipe2` on OS `linux`
2399    fn test_drop_column_added_after_root() {
2400        let desc = RelationDesc::builder()
2401            .with_column("a", SqlScalarType::Bool.nullable(true))
2402            .finish();
2403        let mut versioned = VersionedRelationDesc::new(desc);
2404
2405        let v1 = versioned.add_column("b", SqlScalarType::String.nullable(false));
2406        let v2 = versioned.drop_column("b");
2407        assert_eq!(v1, RelationVersion(1));
2408        assert_eq!(v2, RelationVersion(2));
2409
2410        assert_eq!(
2411            versioned
2412                .at_version(RelationVersionSelector::Specific(v1))
2413                .arity(),
2414            2
2415        );
2416        assert_eq!(
2417            versioned
2418                .at_version(RelationVersionSelector::Specific(v2))
2419                .arity(),
2420            1
2421        );
2422    }
2423
2424    #[mz_ore::test]
2425    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `pipe2` on OS `linux`
2426    fn relation_desc_proto_rejects_corrupt_shapes() {
2427        fn proto(num_types: usize, num_names: usize, keys: Vec<Vec<usize>>) -> ProtoRelationDesc {
2428            let mut typ = SqlRelationType::new(vec![SqlScalarType::Bool.nullable(true); num_types]);
2429            typ.keys = keys;
2430            ProtoRelationDesc {
2431                typ: Some(typ.into_proto()),
2432                names: (0..num_names)
2433                    .map(|i| ColumnName::from(format!("c{i}")).into_proto())
2434                    .collect(),
2435                metadata: vec![],
2436            }
2437        }
2438
2439        // A desc with more names than types panics in `iter()`, one with more types than names
2440        // panics in `into_iter()`, and an out of bounds key is what `validate` calls corruption.
2441        // All three re-encode identically, so a proto round-trip oracle cannot see them.
2442        for (num_types, num_names) in [(0, 1), (2, 0), (3, 1)] {
2443            let err = RelationDesc::from_proto(proto(num_types, num_names, vec![]))
2444                .expect_err("length mismatch must be rejected");
2445            assert!(err.to_string().contains("length mismatch"), "{err}");
2446        }
2447        let err = RelationDesc::from_proto(proto(1, 1, vec![vec![7]]))
2448            .expect_err("out of bounds key must be rejected");
2449        assert!(err.to_string().contains("out of bounds"), "{err}");
2450
2451        // The well formed shape still decodes, and stays usable.
2452        let desc = RelationDesc::from_proto(proto(2, 2, vec![vec![1]])).expect("valid");
2453        assert_eq!(desc.iter().count(), 2);
2454    }
2455
2456    #[mz_ore::test]
2457    #[cfg_attr(miri, ignore)]
2458    fn apply_demand() {
2459        let desc = RelationDesc::builder()
2460            .with_column("a", SqlScalarType::String.nullable(true))
2461            .with_column("b", SqlScalarType::Int64.nullable(false))
2462            .with_column("c", SqlScalarType::Time.nullable(false))
2463            .finish();
2464        let desc = desc.apply_demand(&BTreeSet::from([0, 2]));
2465        assert_eq!(desc.arity(), 2);
2466        // TODO(parkmycar): Move validate onto RelationDesc.
2467        VersionedRelationDesc::new(desc).validate();
2468    }
2469
2470    #[mz_ore::test]
2471    #[cfg_attr(miri, ignore)]
2472    fn smoketest_column_index_stable_ident() {
2473        let idx_a = ColumnIndex(42);
2474        // Note(parkmycar): This should never change.
2475        assert_eq!(idx_a.to_stable_name(), "42");
2476    }
2477
2478    #[mz_ore::test]
2479    #[cfg_attr(miri, ignore)] // too slow
2480    fn proptest_relation_desc_roundtrips() {
2481        fn testcase(og: RelationDesc) {
2482            let bytes = og.into_proto().encode_to_vec();
2483            let proto = ProtoRelationDesc::decode(&bytes[..]).unwrap();
2484            let rnd = RelationDesc::from_proto(proto).unwrap();
2485
2486            assert_eq!(og, rnd);
2487        }
2488
2489        proptest!(|(desc in any::<RelationDesc>())| {
2490            testcase(desc);
2491        });
2492
2493        let strat = any::<RelationDesc>().prop_flat_map(|desc| {
2494            arb_relation_desc_diff(&desc).prop_map(move |diffs| (desc.clone(), diffs))
2495        });
2496
2497        proptest!(|((mut desc, diffs) in strat)| {
2498            for diff in diffs {
2499                diff.apply(&mut desc);
2500            };
2501            testcase(desc);
2502        });
2503    }
2504}