1use 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#[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 pub scalar_type: SqlScalarType,
64 #[serde(default = "return_true")]
66 pub nullable: bool,
67}
68
69#[inline(always)]
75fn return_true() -> bool {
76 true
77}
78
79impl SqlColumnType {
80 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 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 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 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 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 soft_panic_or_log!("repr type error: sql_union({self:?}, {other:?}): {e}");
143 Ok(SqlColumnType::from_repr(&typ))
144 }
145 Err(_) => {
146 Err(e)
149 }
150 }
151 })
152 }
153
154 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 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#[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 pub column_types: Vec<SqlColumnType>,
213 #[serde(default)]
223 pub keys: Vec<Vec<usize>>,
224}
225
226impl SqlRelationType {
227 pub fn empty() -> Self {
230 SqlRelationType::new(vec![])
231 }
232
233 pub fn new(column_types: Vec<SqlColumnType>) -> Self {
237 SqlRelationType {
238 column_types,
239 keys: Vec::new(),
240 }
241 }
242
243 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 pub fn arity(&self) -> usize {
261 self.column_types.len()
262 }
263
264 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 pub fn columns(&self) -> &[SqlColumnType] {
279 &self.column_types
280 }
281
282 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 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#[derive(
347 Clone,
348 Debug,
349 Eq,
350 PartialEq,
351 Ord,
352 PartialOrd,
353 Serialize,
354 Deserialize,
355 Hash
356)]
357pub struct ReprRelationType {
358 pub column_types: Vec<ReprColumnType>,
360 #[serde(default)]
370 pub keys: Vec<Vec<usize>>,
371}
372
373impl ReprRelationType {
374 pub fn empty() -> Self {
377 ReprRelationType::new(vec![])
378 }
379
380 pub fn new(column_types: Vec<ReprColumnType>) -> Self {
384 ReprRelationType {
385 column_types,
386 keys: Vec::new(),
387 }
388 }
389
390 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 pub fn arity(&self) -> usize {
408 self.column_types.len()
409 }
410
411 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 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 pub scalar_type: ReprScalarType,
457 #[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 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 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#[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 #[inline(always)]
536 pub fn as_str(&self) -> &str {
537 &*self
538 }
539
540 pub fn as_mut_boxed_str(&mut self) -> &mut Box<str> {
542 &mut self.0
543 }
544
545 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 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 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 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
646pub const UNKNOWN_COLUMN_NAME: &str = "?column?";
648
649#[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 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#[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 pub fn root() -> Self {
701 RelationVersion(0)
702 }
703
704 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 pub fn into_raw(self) -> u64 {
717 self.0
718 }
719
720 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#[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#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash)]
832struct ColumnMetadata {
833 name: ColumnName,
835 typ_idx: usize,
837 added: RelationVersion,
839 dropped: Option<RelationVersion>,
841}
842
843#[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 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 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 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 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 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 pub fn builder() -> RelationDescBuilder {
1033 RelationDescBuilder::default()
1034 }
1035
1036 pub fn empty() -> Self {
1039 RelationDesc {
1040 typ: SqlRelationType::empty(),
1041 metadata: BTreeMap::default(),
1042 }
1043 }
1044
1045 pub fn is_empty(&self) -> bool {
1047 self == &Self::empty()
1048 }
1049
1050 pub fn len(&self) -> usize {
1052 self.typ().column_types.len()
1053 }
1054
1055 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 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 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 pub fn with_key(mut self, indices: Vec<usize>) -> Self {
1140 self.typ = self.typ.with_key(indices);
1141 self
1142 }
1143
1144 pub fn without_keys(mut self) -> Self {
1146 self.typ.keys.clear();
1147 self
1148 }
1149
1150 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 pub fn arity(&self) -> usize {
1167 self.typ.arity()
1168 }
1169
1170 pub fn typ(&self) -> &SqlRelationType {
1172 &self.typ
1173 }
1174
1175 pub fn into_typ(self) -> SqlRelationType {
1177 self.typ
1178 }
1179
1180 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 pub fn iter_types(&self) -> impl Iterator<Item = &SqlColumnType> {
1190 self.typ.column_types.iter()
1191 }
1192
1193 pub fn iter_names(&self) -> impl Iterator<Item = &ColumnName> {
1195 self.metadata.values().map(|meta| &meta.name)
1196 }
1197
1198 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 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 pub fn contains_index(&self, idx: &ColumnIndex) -> bool {
1217 self.metadata.contains_key(idx)
1218 }
1219
1220 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 pub fn get_name(&self, i: usize) -> &ColumnName {
1239 self.get_name_idx(&ColumnIndex(i))
1241 }
1242
1243 pub fn get_name_idx(&self, idx: &ColumnIndex) -> &ColumnName {
1249 &self.metadata.get(idx).expect("should exist").name
1250 }
1251
1252 pub fn get_name_mut(&mut self, i: usize) -> &mut ColumnName {
1258 &mut self
1260 .metadata
1261 .get_mut(&ColumnIndex(i))
1262 .expect("should exist")
1263 .name
1264 }
1265
1266 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 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 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 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 pub fn apply_demand(&self, demands: &BTreeSet<usize>) -> RelationDesc {
1402 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 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 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#[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#[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#[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#[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#[derive(Debug, Clone, PartialEq, Eq)]
1527pub struct RelationDescDiff {
1528 pub column_diffs: BTreeMap<usize, ColumnDiff>,
1530 pub key_diff: Option<KeyDiff>,
1532}
1533
1534impl RelationDescDiff {
1535 pub fn is_empty(&self) -> bool {
1537 self.column_diffs.is_empty() && self.key_diff.is_none()
1538 }
1539}
1540
1541#[derive(Debug, Clone, PartialEq, Eq)]
1543pub enum ColumnDiff {
1544 Missing { name: ColumnName },
1546 Extra { name: ColumnName },
1548 TypeMismatch {
1550 name: ColumnName,
1551 left: SqlScalarType,
1552 right: SqlScalarType,
1553 },
1554 NullabilityMismatch {
1556 name: ColumnName,
1557 left: bool,
1558 right: bool,
1559 },
1560 NameMismatch { left: ColumnName, right: ColumnName },
1562}
1563
1564#[derive(Debug, Clone, PartialEq, Eq)]
1566pub struct KeyDiff {
1567 pub left: BTreeSet<Vec<ColumnName>>,
1569 pub right: BTreeSet<Vec<ColumnName>>,
1571}
1572
1573#[derive(Clone, Default, Debug, PartialEq, Eq)]
1575pub struct RelationDescBuilder {
1576 columns: Vec<(ColumnName, SqlColumnType)>,
1578 keys: Vec<Vec<usize>>,
1580}
1581
1582impl RelationDescBuilder {
1583 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 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 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 pub fn without_keys(mut self) -> RelationDescBuilder {
1617 self.keys.clear();
1618 assert_eq!(self.keys.len(), 0);
1619 self
1620 }
1621
1622 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 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#[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#[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 #[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 #[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 assert_none!(col.dropped, "column was already dropped");
1740 col.dropped = Some(new_version);
1741
1742 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 pub fn latest(&self) -> RelationDesc {
1757 self.inner.clone()
1758 }
1759
1760 pub fn at_version(&self, version: RelationVersionSelector) -> RelationDesc {
1762 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 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 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 .map(|meta| meta.dropped.unwrap_or(meta.added))
1834 .max()
1835 .unwrap_or_else(RelationVersion::root)
1837 }
1838
1839 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 let versions = desc
1876 .metadata
1877 .values()
1878 .flat_map(|meta| [Some(meta.added), meta.dropped])
1879 .flatten()
1880 .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 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#[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#[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 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 #[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)] 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 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 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)] 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 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 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)] 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)] 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)] 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)] 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 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 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 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 assert_eq!(idx_a.to_stable_name(), "42");
2476 }
2477
2478 #[mz_ore::test]
2479 #[cfg_attr(miri, ignore)] 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}