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
606#[cfg(any(test, feature = "proptest"))]
607impl proptest::arbitrary::Arbitrary for ColumnName {
608 type Parameters = ();
609 type Strategy = BoxedStrategy<ColumnName>;
610
611 fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
612 let mut weights = vec![(50, Just(1..8)), (20, Just(8..16))];
615 if std::env::var("PROPTEST_LARGE_DATA").is_ok() {
616 weights.extend([
617 (5, Just(16..128)),
618 (1, Just(128..1024)),
619 (1, Just(1024..4096)),
620 ]);
621 }
622 let name_length = Union::new_weighted(weights);
623
624 let char_strat = Rc::new(Union::new_weighted(vec![
627 (50, proptest::char::range('A', 'z').boxed()),
628 (1, any::<char>().boxed()),
629 ]));
630
631 name_length
632 .prop_flat_map(move |length| proptest::collection::vec(Rc::clone(&char_strat), length))
633 .prop_map(|chars| ColumnName(chars.into_iter().collect::<Box<str>>()))
634 .no_shrink()
635 .boxed()
636 }
637}
638
639pub const UNKNOWN_COLUMN_NAME: &str = "?column?";
641
642pub static ORDINALITY_COL_NAME: &str = "ordinality";
644
645#[derive(
647 Clone,
648 Copy,
649 Debug,
650 Eq,
651 PartialEq,
652 PartialOrd,
653 Ord,
654 Serialize,
655 Deserialize,
656 Hash
657)]
658pub struct ColumnIndex(usize);
659
660#[cfg(any(test, feature = "proptest"))]
661static_assertions::assert_not_impl_all!(ColumnIndex: Arbitrary);
662
663impl ColumnIndex {
664 pub fn to_stable_name(&self) -> String {
666 self.0.to_string()
667 }
668
669 pub fn to_raw(&self) -> usize {
670 self.0
671 }
672
673 pub fn from_raw(val: usize) -> Self {
674 ColumnIndex(val)
675 }
676}
677
678#[derive(
680 Clone,
681 Copy,
682 Debug,
683 Eq,
684 PartialEq,
685 PartialOrd,
686 Ord,
687 Serialize,
688 Deserialize,
689 Hash
690)]
691#[cfg_attr(any(test, feature = "proptest"), derive(Arbitrary))]
692pub struct RelationVersion(u64);
693
694impl RelationVersion {
695 pub fn root() -> Self {
697 RelationVersion(0)
698 }
699
700 pub fn bump(&self) -> Self {
702 let next_version = self
703 .0
704 .checked_add(1)
705 .expect("added more than u64::MAX columns?");
706 RelationVersion(next_version)
707 }
708
709 pub fn into_raw(self) -> u64 {
713 self.0
714 }
715
716 pub fn from_raw(val: u64) -> RelationVersion {
720 RelationVersion(val)
721 }
722}
723
724impl From<RelationVersion> for SchemaId {
725 fn from(value: RelationVersion) -> Self {
726 SchemaId(usize::cast_from(value.0))
727 }
728}
729
730impl fmt::Display for RelationVersion {
731 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
732 write!(f, "v{}", self.0)
733 }
734}
735
736impl RustType<ProtoRelationVersion> for RelationVersion {
737 fn into_proto(&self) -> ProtoRelationVersion {
738 ProtoRelationVersion { value: self.0 }
739 }
740
741 fn from_proto(proto: ProtoRelationVersion) -> Result<Self, TryFromProtoError> {
742 Ok(RelationVersion(proto.value))
743 }
744}
745
746#[derive(
753 Clone,
754 Copy,
755 Debug,
756 PartialEq,
757 Eq,
758 PartialOrd,
759 Ord,
760 Hash,
761 serde::Serialize
762)]
763pub enum SemanticType {
764 CatalogItemId,
765 GlobalId,
766 ClusterId,
767 ReplicaId,
768 SchemaId,
769 DatabaseId,
770 RoleId,
771 NetworkPolicyId,
772 ShardId,
773 OID,
774 ObjectType,
775 ConnectionType,
776 SourceType,
777 MzTimestamp,
778 WallclockTimestamp,
779 ByteCount,
780 RecordCount,
781 CreditRate,
782 SqlDefinition,
783 RedactedSqlDefinition,
784}
785
786impl fmt::Display for SemanticType {
787 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
788 let s = match self {
789 SemanticType::CatalogItemId => "CatalogItemId",
790 SemanticType::GlobalId => "GlobalId",
791 SemanticType::ClusterId => "ClusterId",
792 SemanticType::ReplicaId => "ReplicaId",
793 SemanticType::SchemaId => "SchemaId",
794 SemanticType::DatabaseId => "DatabaseId",
795 SemanticType::RoleId => "RoleId",
796 SemanticType::NetworkPolicyId => "NetworkPolicyId",
797 SemanticType::ShardId => "ShardId",
798 SemanticType::OID => "OID",
799 SemanticType::ObjectType => "ObjectType",
800 SemanticType::ConnectionType => "ConnectionType",
801 SemanticType::SourceType => "SourceType",
802 SemanticType::MzTimestamp => "MzTimestamp",
803 SemanticType::WallclockTimestamp => "WallclockTimestamp",
804 SemanticType::ByteCount => "ByteCount",
805 SemanticType::RecordCount => "RecordCount",
806 SemanticType::CreditRate => "CreditRate",
807 SemanticType::SqlDefinition => "SqlDefinition",
808 SemanticType::RedactedSqlDefinition => "RedactedSqlDefinition",
809 };
810 f.write_str(s)
811 }
812}
813
814#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, Hash)]
816struct ColumnMetadata {
817 name: ColumnName,
819 typ_idx: usize,
821 added: RelationVersion,
823 dropped: Option<RelationVersion>,
825}
826
827#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
895pub struct RelationDesc {
896 typ: SqlRelationType,
897 metadata: BTreeMap<ColumnIndex, ColumnMetadata>,
898}
899
900impl RustType<ProtoRelationDesc> for RelationDesc {
901 fn into_proto(&self) -> ProtoRelationDesc {
910 let (names, metadata): (Vec<_>, Vec<_>) = self
911 .metadata
912 .values()
913 .map(|meta| {
914 let metadata = ProtoColumnMetadata {
915 added: Some(meta.added.into_proto()),
916 dropped: meta.dropped.map(|v| v.into_proto()),
917 };
918 (meta.name.into_proto(), metadata)
919 })
920 .unzip();
921
922 let is_all_default_metadata = metadata.iter().all(|meta| {
928 meta.added == Some(RelationVersion::root().into_proto()) && meta.dropped == None
929 });
930 let metadata = if is_all_default_metadata {
931 Vec::new()
932 } else {
933 metadata
934 };
935
936 ProtoRelationDesc {
937 typ: Some(self.typ.into_proto()),
938 names,
939 metadata,
940 }
941 }
942
943 fn from_proto(proto: ProtoRelationDesc) -> Result<Self, TryFromProtoError> {
944 let typ: SqlRelationType = proto.typ.into_rust_if_some("ProtoRelationDesc::typ")?;
945
946 if proto.names.len() != typ.column_types.len() {
951 return Err(TryFromProtoError::InvalidFieldError(format!(
952 "ProtoRelationDesc: names ({}) and column_types ({}) length mismatch",
953 proto.names.len(),
954 typ.column_types.len()
955 )));
956 }
957 if let Some(key) = typ
958 .keys
959 .iter()
960 .flatten()
961 .find(|key| **key >= typ.column_types.len())
962 {
963 return Err(TryFromProtoError::InvalidFieldError(format!(
964 "ProtoRelationDesc: key index {key} out of bounds for {} columns",
965 typ.column_types.len()
966 )));
967 }
968
969 let proto_metadata: Box<dyn Iterator<Item = _>> = if proto.metadata.is_empty() {
975 let val = ProtoColumnMetadata {
976 added: Some(RelationVersion::root().into_proto()),
977 dropped: None,
978 };
979 Box::new(itertools::repeat_n(val, proto.names.len()))
980 } else {
981 if proto.names.len() != proto.metadata.len() {
985 return Err(TryFromProtoError::InvalidFieldError(format!(
986 "ProtoRelationDesc: names ({}) and metadata ({}) length mismatch",
987 proto.names.len(),
988 proto.metadata.len()
989 )));
990 }
991 Box::new(proto.metadata.into_iter())
992 };
993
994 let metadata = proto
995 .names
996 .into_iter()
997 .zip_eq(proto_metadata)
998 .enumerate()
999 .map(|(idx, (name, metadata))| {
1000 let meta = ColumnMetadata {
1001 name: name.into_rust()?,
1002 typ_idx: idx,
1003 added: metadata.added.into_rust_if_some("ColumnMetadata::added")?,
1004 dropped: metadata.dropped.into_rust()?,
1005 };
1006 Ok::<_, TryFromProtoError>((ColumnIndex(idx), meta))
1007 })
1008 .collect::<Result<_, _>>()?;
1009
1010 Ok(RelationDesc { typ, metadata })
1011 }
1012}
1013
1014impl RelationDesc {
1015 pub fn builder() -> RelationDescBuilder {
1017 RelationDescBuilder::default()
1018 }
1019
1020 pub fn empty() -> Self {
1023 RelationDesc {
1024 typ: SqlRelationType::empty(),
1025 metadata: BTreeMap::default(),
1026 }
1027 }
1028
1029 pub fn is_empty(&self) -> bool {
1031 self == &Self::empty()
1032 }
1033
1034 pub fn len(&self) -> usize {
1036 self.typ().column_types.len()
1037 }
1038
1039 pub fn new<I, N>(typ: SqlRelationType, names: I) -> Self
1047 where
1048 I: IntoIterator<Item = N>,
1049 N: Into<ColumnName>,
1050 {
1051 let metadata: BTreeMap<_, _> = names
1052 .into_iter()
1053 .enumerate()
1054 .map(|(idx, name)| {
1055 let col_idx = ColumnIndex(idx);
1056 let metadata = ColumnMetadata {
1057 name: name.into(),
1058 typ_idx: idx,
1059 added: RelationVersion::root(),
1060 dropped: None,
1061 };
1062 (col_idx, metadata)
1063 })
1064 .collect();
1065
1066 assert_eq!(typ.column_types.len(), metadata.len());
1068
1069 RelationDesc { typ, metadata }
1070 }
1071
1072 pub fn from_names_and_types<I, T, N>(iter: I) -> Self
1073 where
1074 I: IntoIterator<Item = (N, T)>,
1075 T: Into<SqlColumnType>,
1076 N: Into<ColumnName>,
1077 {
1078 let (names, types): (Vec<_>, Vec<_>) = iter.into_iter().unzip();
1079 let types = types.into_iter().map(Into::into).collect();
1080 let typ = SqlRelationType::new(types);
1081 Self::new(typ, names)
1082 }
1083
1084 pub fn concat(mut self, other: Self) -> Self {
1094 let self_len = self.typ.column_types.len();
1095
1096 for (typ, (_col_idx, meta)) in other.typ.column_types.into_iter().zip_eq(other.metadata) {
1097 assert_eq!(meta.added, RelationVersion::root());
1098 assert_none!(meta.dropped);
1099
1100 let new_idx = self.typ.columns().len();
1101 let new_meta = ColumnMetadata {
1102 name: meta.name,
1103 typ_idx: new_idx,
1104 added: RelationVersion::root(),
1105 dropped: None,
1106 };
1107
1108 self.typ.column_types.push(typ);
1109 let prev = self.metadata.insert(ColumnIndex(new_idx), new_meta);
1110
1111 assert_eq!(self.metadata.len(), self.typ.columns().len());
1112 assert_none!(prev);
1113 }
1114
1115 for k in other.typ.keys {
1116 let k = k.into_iter().map(|idx| idx + self_len).collect();
1117 self = self.with_key(k);
1118 }
1119 self
1120 }
1121
1122 pub fn with_key(mut self, indices: Vec<usize>) -> Self {
1124 self.typ = self.typ.with_key(indices);
1125 self
1126 }
1127
1128 pub fn without_keys(mut self) -> Self {
1130 self.typ.keys.clear();
1131 self
1132 }
1133
1134 pub fn with_names<I, N>(self, names: I) -> Self
1142 where
1143 I: IntoIterator<Item = N>,
1144 N: Into<ColumnName>,
1145 {
1146 Self::new(self.typ, names)
1147 }
1148
1149 pub fn arity(&self) -> usize {
1151 self.typ.arity()
1152 }
1153
1154 pub fn typ(&self) -> &SqlRelationType {
1156 &self.typ
1157 }
1158
1159 pub fn into_typ(self) -> SqlRelationType {
1161 self.typ
1162 }
1163
1164 pub fn iter(&self) -> impl Iterator<Item = (&ColumnName, &SqlColumnType)> {
1166 self.metadata.values().map(|meta| {
1167 let typ = &self.typ.columns()[meta.typ_idx];
1168 (&meta.name, typ)
1169 })
1170 }
1171
1172 pub fn iter_types(&self) -> impl Iterator<Item = &SqlColumnType> {
1174 self.typ.column_types.iter()
1175 }
1176
1177 pub fn iter_names(&self) -> impl Iterator<Item = &ColumnName> {
1179 self.metadata.values().map(|meta| &meta.name)
1180 }
1181
1182 pub fn iter_all(&self) -> impl Iterator<Item = (&ColumnIndex, &ColumnName, &SqlColumnType)> {
1184 self.metadata.iter().map(|(col_idx, metadata)| {
1185 let col_typ = &self.typ.columns()[metadata.typ_idx];
1186 (col_idx, &metadata.name, col_typ)
1187 })
1188 }
1189
1190 pub fn iter_similar_names<'a>(
1193 &'a self,
1194 name: &'a ColumnName,
1195 ) -> impl Iterator<Item = &'a ColumnName> {
1196 self.iter_names().filter(|n| n.is_similar(name))
1197 }
1198
1199 pub fn contains_index(&self, idx: &ColumnIndex) -> bool {
1201 self.metadata.contains_key(idx)
1202 }
1203
1204 pub fn get_by_name(&self, name: &ColumnName) -> Option<(usize, &SqlColumnType)> {
1210 self.iter_names()
1211 .position(|n| n == name)
1212 .map(|i| (i, &self.typ.column_types[i]))
1213 }
1214
1215 pub fn get_name(&self, i: usize) -> &ColumnName {
1223 self.get_name_idx(&ColumnIndex(i))
1225 }
1226
1227 pub fn get_name_idx(&self, idx: &ColumnIndex) -> &ColumnName {
1233 &self.metadata.get(idx).expect("should exist").name
1234 }
1235
1236 pub fn get_name_mut(&mut self, i: usize) -> &mut ColumnName {
1242 &mut self
1244 .metadata
1245 .get_mut(&ColumnIndex(i))
1246 .expect("should exist")
1247 .name
1248 }
1249
1250 pub fn get_type(&self, idx: &ColumnIndex) -> &SqlColumnType {
1256 let typ_idx = self.metadata.get(idx).expect("should exist").typ_idx;
1257 &self.typ.column_types[typ_idx]
1258 }
1259
1260 pub fn get_unambiguous_name(&self, i: usize) -> Option<&ColumnName> {
1269 let name = self.get_name(i);
1270 if self.iter_names().filter(|n| *n == name).count() == 1 {
1271 Some(name)
1272 } else {
1273 None
1274 }
1275 }
1276
1277 pub fn constraints_met(&self, i: usize, d: &Datum) -> Result<(), NotNullViolation> {
1282 let name = self.get_name(i);
1283 let typ = &self.typ.column_types[i];
1284 if d == &Datum::Null && !typ.nullable {
1285 Err(NotNullViolation(name.clone()))
1286 } else {
1287 Ok(())
1288 }
1289 }
1290
1291 pub fn diff(&self, other: &RelationDesc) -> RelationDescDiff {
1306 assert_eq!(self.metadata.len(), self.typ.columns().len());
1307 assert_eq!(other.metadata.len(), other.typ.columns().len());
1308 for (idx, meta) in self.metadata.iter().chain(other.metadata.iter()) {
1309 assert_eq!(meta.typ_idx, idx.0);
1310 assert_eq!(meta.added, RelationVersion::root());
1311 assert_none!(meta.dropped);
1312 }
1313
1314 let mut column_diffs = BTreeMap::new();
1315 let mut key_diff = None;
1316
1317 let left_arity = self.arity();
1318 let right_arity = other.arity();
1319 let common_arity = std::cmp::min(left_arity, right_arity);
1320
1321 for idx in 0..common_arity {
1322 let left_name = self.get_name(idx);
1323 let right_name = other.get_name(idx);
1324 let left_type = &self.typ.column_types[idx];
1325 let right_type = &other.typ.column_types[idx];
1326
1327 if left_name != right_name {
1328 let diff = ColumnDiff::NameMismatch {
1329 left: left_name.clone(),
1330 right: right_name.clone(),
1331 };
1332 column_diffs.insert(idx, diff);
1333 } else if left_type.scalar_type != right_type.scalar_type {
1334 let diff = ColumnDiff::TypeMismatch {
1335 name: left_name.clone(),
1336 left: left_type.scalar_type.clone(),
1337 right: right_type.scalar_type.clone(),
1338 };
1339 column_diffs.insert(idx, diff);
1340 } else if left_type.nullable != right_type.nullable {
1341 let diff = ColumnDiff::NullabilityMismatch {
1342 name: left_name.clone(),
1343 left: left_type.nullable,
1344 right: right_type.nullable,
1345 };
1346 column_diffs.insert(idx, diff);
1347 }
1348 }
1349
1350 for idx in common_arity..left_arity {
1351 let diff = ColumnDiff::Missing {
1352 name: self.get_name(idx).clone(),
1353 };
1354 column_diffs.insert(idx, diff);
1355 }
1356
1357 for idx in common_arity..right_arity {
1358 let diff = ColumnDiff::Extra {
1359 name: other.get_name(idx).clone(),
1360 };
1361 column_diffs.insert(idx, diff);
1362 }
1363
1364 let left_keys: BTreeSet<_> = self.typ.keys.iter().collect();
1365 let right_keys: BTreeSet<_> = other.typ.keys.iter().collect();
1366 if left_keys != right_keys {
1367 let column_names = |desc: &RelationDesc, keys: BTreeSet<&Vec<usize>>| {
1368 keys.iter()
1369 .map(|key| key.iter().map(|&idx| desc.get_name(idx).clone()).collect())
1370 .collect()
1371 };
1372 key_diff = Some(KeyDiff {
1373 left: column_names(self, left_keys),
1374 right: column_names(other, right_keys),
1375 });
1376 }
1377
1378 RelationDescDiff {
1379 column_diffs,
1380 key_diff,
1381 }
1382 }
1383
1384 pub fn apply_demand(&self, demands: &BTreeSet<usize>) -> RelationDesc {
1386 debug_assert!(
1392 self.metadata
1393 .iter()
1394 .enumerate()
1395 .all(|(pos, (idx, meta))| idx.0 == pos && meta.typ_idx == pos),
1396 "apply_demand requires a dense RelationDesc (ColumnIndex == typ_idx): {:?}",
1397 self.metadata,
1398 );
1399 let mut new_desc = self.clone();
1400
1401 let mut removed = 0;
1403 new_desc.metadata.retain(|idx, metadata| {
1404 let retain = demands.contains(&idx.0);
1405 if !retain {
1406 removed += 1;
1407 } else {
1408 metadata.typ_idx -= removed;
1409 }
1410 retain
1411 });
1412
1413 let mut idx = 0;
1415 new_desc.typ.column_types.retain(|_| {
1416 let keep = demands.contains(&idx);
1417 idx += 1;
1418 keep
1419 });
1420
1421 new_desc
1422 }
1423}
1424
1425#[cfg(any(test, feature = "proptest"))]
1426impl Arbitrary for RelationDesc {
1427 type Parameters = ();
1428 type Strategy = BoxedStrategy<RelationDesc>;
1429
1430 fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
1431 let mut weights = vec![(100, Just(0..4)), (50, Just(4..8)), (25, Just(8..16))];
1432 if std::env::var("PROPTEST_LARGE_DATA").is_ok() {
1433 weights.extend([
1434 (12, Just(16..32)),
1435 (6, Just(32..64)),
1436 (3, Just(64..128)),
1437 (1, Just(128..256)),
1438 ]);
1439 }
1440 let num_columns = Union::new_weighted(weights);
1441
1442 num_columns.prop_flat_map(arb_relation_desc).boxed()
1443 }
1444}
1445
1446#[cfg(any(test, feature = "proptest"))]
1449pub fn arb_relation_desc(num_cols: std::ops::Range<usize>) -> impl Strategy<Value = RelationDesc> {
1450 proptest::collection::btree_map(any::<ColumnName>(), any::<SqlColumnType>(), num_cols)
1451 .prop_map(RelationDesc::from_names_and_types)
1452}
1453
1454#[cfg(any(test, feature = "proptest"))]
1456pub fn arb_relation_desc_projection(desc: RelationDesc) -> impl Strategy<Value = RelationDesc> {
1457 let mask: Vec<_> = (0..desc.len()).map(|_| any::<bool>()).collect();
1458 mask.prop_map(move |mask| {
1459 let demands: BTreeSet<_> = mask
1460 .into_iter()
1461 .enumerate()
1462 .filter_map(|(idx, keep)| keep.then_some(idx))
1463 .collect();
1464 desc.apply_demand(&demands)
1465 })
1466}
1467
1468impl IntoIterator for RelationDesc {
1469 type Item = (ColumnName, SqlColumnType);
1470 type IntoIter = Box<dyn Iterator<Item = (ColumnName, SqlColumnType)>>;
1471
1472 fn into_iter(self) -> Self::IntoIter {
1473 let iter = self
1474 .metadata
1475 .into_values()
1476 .zip_eq(self.typ.column_types)
1477 .map(|(meta, typ)| (meta.name, typ));
1478 Box::new(iter)
1479 }
1480}
1481
1482#[cfg(any(test, feature = "proptest"))]
1484pub fn arb_row_for_relation(desc: &RelationDesc) -> impl Strategy<Value = Row> + use<> {
1485 let datums: Vec<_> = desc
1486 .typ()
1487 .columns()
1488 .iter()
1489 .cloned()
1490 .map(arb_datum_for_column)
1491 .collect();
1492 datums.prop_map(|x| Row::pack(x.iter().map(Datum::from)))
1493}
1494
1495#[derive(Debug, PartialEq, Eq)]
1497pub struct NotNullViolation(pub ColumnName);
1498
1499impl fmt::Display for NotNullViolation {
1500 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1501 write!(
1502 f,
1503 "null value in column {} violates not-null constraint",
1504 self.0.quoted()
1505 )
1506 }
1507}
1508
1509#[derive(Debug, Clone, PartialEq, Eq)]
1511pub struct RelationDescDiff {
1512 pub column_diffs: BTreeMap<usize, ColumnDiff>,
1514 pub key_diff: Option<KeyDiff>,
1516}
1517
1518impl RelationDescDiff {
1519 pub fn is_empty(&self) -> bool {
1521 self.column_diffs.is_empty() && self.key_diff.is_none()
1522 }
1523}
1524
1525#[derive(Debug, Clone, PartialEq, Eq)]
1527pub enum ColumnDiff {
1528 Missing { name: ColumnName },
1530 Extra { name: ColumnName },
1532 TypeMismatch {
1534 name: ColumnName,
1535 left: SqlScalarType,
1536 right: SqlScalarType,
1537 },
1538 NullabilityMismatch {
1540 name: ColumnName,
1541 left: bool,
1542 right: bool,
1543 },
1544 NameMismatch { left: ColumnName, right: ColumnName },
1546}
1547
1548#[derive(Debug, Clone, PartialEq, Eq)]
1550pub struct KeyDiff {
1551 pub left: BTreeSet<Vec<ColumnName>>,
1553 pub right: BTreeSet<Vec<ColumnName>>,
1555}
1556
1557#[derive(Clone, Default, Debug, PartialEq, Eq)]
1559pub struct RelationDescBuilder {
1560 columns: Vec<(ColumnName, SqlColumnType)>,
1562 keys: Vec<Vec<usize>>,
1564}
1565
1566impl RelationDescBuilder {
1567 pub fn with_column<N: Into<ColumnName>>(
1569 mut self,
1570 name: N,
1571 ty: SqlColumnType,
1572 ) -> RelationDescBuilder {
1573 let name = name.into();
1574 self.columns.push((name, ty));
1575 self
1576 }
1577
1578 pub fn with_columns<I, T, N>(mut self, iter: I) -> Self
1580 where
1581 I: IntoIterator<Item = (N, T)>,
1582 T: Into<SqlColumnType>,
1583 N: Into<ColumnName>,
1584 {
1585 self.columns
1586 .extend(iter.into_iter().map(|(name, ty)| (name.into(), ty.into())));
1587 self
1588 }
1589
1590 pub fn with_key(mut self, mut indices: Vec<usize>) -> RelationDescBuilder {
1592 indices.sort_unstable();
1593 if !self.keys.contains(&indices) {
1594 self.keys.push(indices);
1595 }
1596 self
1597 }
1598
1599 pub fn without_keys(mut self) -> RelationDescBuilder {
1601 self.keys.clear();
1602 assert_eq!(self.keys.len(), 0);
1603 self
1604 }
1605
1606 pub fn concat(mut self, other: Self) -> Self {
1608 let self_len = self.columns.len();
1609
1610 self.columns.extend(other.columns);
1611 for k in other.keys {
1612 let k = k.into_iter().map(|idx| idx + self_len).collect();
1613 self = self.with_key(k);
1614 }
1615
1616 self
1617 }
1618
1619 pub fn finish(self) -> RelationDesc {
1621 let mut desc = RelationDesc::from_names_and_types(self.columns);
1622 desc.typ = desc.typ.with_keys(self.keys);
1623 desc
1624 }
1625}
1626
1627#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize)]
1629pub enum RelationVersionSelector {
1630 Specific(RelationVersion),
1631 Latest,
1632}
1633
1634impl RelationVersionSelector {
1635 pub fn specific(version: u64) -> Self {
1636 RelationVersionSelector::Specific(RelationVersion(version))
1637 }
1638}
1639
1640#[derive(Debug, Clone, Serialize)]
1646pub struct VersionedRelationDesc {
1647 inner: RelationDesc,
1648}
1649
1650impl VersionedRelationDesc {
1651 pub fn new(inner: RelationDesc) -> Self {
1652 VersionedRelationDesc { inner }
1653 }
1654
1655 #[must_use]
1663 pub fn add_column<N, T>(&mut self, name: N, typ: T) -> RelationVersion
1664 where
1665 N: Into<ColumnName>,
1666 T: Into<SqlColumnType>,
1667 {
1668 let latest_version = self.latest_version();
1669 let new_version = latest_version.bump();
1670
1671 let name = name.into();
1672 let existing = self
1673 .inner
1674 .metadata
1675 .iter()
1676 .find(|(_, meta)| meta.name == name && meta.dropped.is_none());
1677 if let Some(existing) = existing {
1678 panic!("column named '{name}' already exists! {existing:?}");
1679 }
1680
1681 let next_idx = self.inner.metadata.len();
1682 let col_meta = ColumnMetadata {
1683 name,
1684 typ_idx: next_idx,
1685 added: new_version,
1686 dropped: None,
1687 };
1688
1689 self.inner.typ.column_types.push(typ.into());
1690 let prev = self.inner.metadata.insert(ColumnIndex(next_idx), col_meta);
1691
1692 assert_none!(prev, "column index overlap!");
1693 self.validate();
1694
1695 new_version
1696 }
1697
1698 #[must_use]
1707 pub fn drop_column<N>(&mut self, name: N) -> RelationVersion
1708 where
1709 N: Into<ColumnName>,
1710 {
1711 let name = name.into();
1712 let latest_version = self.latest_version();
1713 let new_version = latest_version.bump();
1714
1715 let col = self
1716 .inner
1717 .metadata
1718 .values_mut()
1719 .find(|meta| meta.name == name && meta.dropped.is_none())
1720 .expect("column to exist");
1721
1722 assert_none!(col.dropped, "column was already dropped");
1724 col.dropped = Some(new_version);
1725
1726 let dropped_key = self
1728 .inner
1729 .typ
1730 .keys
1731 .iter()
1732 .any(|keys| keys.contains(&col.typ_idx));
1733 assert!(!dropped_key, "column being dropped was used as a key");
1734
1735 self.validate();
1736 new_version
1737 }
1738
1739 pub fn latest(&self) -> RelationDesc {
1741 self.inner.clone()
1742 }
1743
1744 pub fn at_version(&self, version: RelationVersionSelector) -> RelationDesc {
1746 let up_to_version = match version {
1748 RelationVersionSelector::Latest => RelationVersion(u64::MAX),
1749 RelationVersionSelector::Specific(v) => v,
1750 };
1751
1752 let valid_columns = self.inner.metadata.iter().filter(|(_col_idx, meta)| {
1753 let added = meta.added <= up_to_version;
1754 let dropped = meta
1755 .dropped
1756 .map(|dropped_at| up_to_version >= dropped_at)
1757 .unwrap_or(false);
1758
1759 added && !dropped
1760 });
1761
1762 let mut column_types = Vec::new();
1763 let mut column_metas = BTreeMap::new();
1764
1765 for (col_idx, meta) in valid_columns {
1772 let new_meta = ColumnMetadata {
1773 name: meta.name.clone(),
1774 typ_idx: column_types.len(),
1775 added: meta.added.clone(),
1776 dropped: meta.dropped.clone(),
1777 };
1778 column_types.push(self.inner.typ.columns()[meta.typ_idx].clone());
1779 column_metas.insert(*col_idx, new_meta);
1780 }
1781
1782 let keys = self
1788 .inner
1789 .typ
1790 .keys
1791 .iter()
1792 .map(|keys| {
1793 keys.iter()
1794 .map(|key_idx| {
1795 let metadata = column_metas
1796 .get(&ColumnIndex(*key_idx))
1797 .expect("found key for column that doesn't exist");
1798 metadata.typ_idx
1799 })
1800 .collect()
1801 })
1802 .collect();
1803
1804 let relation_type = SqlRelationType { column_types, keys };
1805
1806 RelationDesc {
1807 typ: relation_type,
1808 metadata: column_metas,
1809 }
1810 }
1811
1812 pub fn latest_version(&self) -> RelationVersion {
1813 self.inner
1814 .metadata
1815 .values()
1816 .map(|meta| meta.dropped.unwrap_or(meta.added))
1818 .max()
1819 .unwrap_or_else(RelationVersion::root)
1821 }
1822
1823 fn validate(&self) {
1829 fn validate_inner(desc: &RelationDesc) -> Result<(), anyhow::Error> {
1830 if desc.typ.column_types.len() != desc.metadata.len() {
1831 anyhow::bail!("mismatch between number of types and metadatas");
1832 }
1833
1834 for (col_idx, meta) in &desc.metadata {
1835 if col_idx.0 > desc.metadata.len() {
1836 anyhow::bail!("column index out of bounds");
1837 }
1838 if meta.added >= meta.dropped.unwrap_or(RelationVersion(u64::MAX)) {
1839 anyhow::bail!("column was added after it was dropped?");
1840 }
1841 if desc.typ().columns().get(meta.typ_idx).is_none() {
1842 anyhow::bail!("typ_idx incorrect");
1843 }
1844 }
1845
1846 for keys in &desc.typ.keys {
1847 for key in keys {
1848 if *key >= desc.typ.column_types.len() {
1849 anyhow::bail!("key index was out of bounds!");
1850 }
1851 }
1852 }
1853
1854 let versions = desc
1860 .metadata
1861 .values()
1862 .flat_map(|meta| [Some(meta.added), meta.dropped])
1863 .flatten()
1864 .filter(|version| *version != RelationVersion::root());
1866 let mut max = 0;
1867 let mut sum = 0;
1868 for version in versions {
1869 max = std::cmp::max(max, version.0);
1870 sum += version.0;
1871 }
1872
1873 if sum != (max * (max + 1) / 2) {
1883 anyhow::bail!("there is a duplicate or missing relation version");
1884 }
1885
1886 Ok(())
1887 }
1888
1889 assert_ok!(validate_inner(&self.inner), "validate failed! {self:?}");
1890 }
1891}
1892
1893#[derive(Debug)]
1896#[cfg(any(test, feature = "proptest"))]
1897pub enum PropRelationDescDiff {
1898 AddColumn {
1899 name: ColumnName,
1900 typ: SqlColumnType,
1901 },
1902 DropColumn {
1903 name: ColumnName,
1904 },
1905 ToggleNullability {
1906 name: ColumnName,
1907 },
1908 ChangeType {
1909 name: ColumnName,
1910 typ: SqlColumnType,
1911 },
1912}
1913
1914#[cfg(any(test, feature = "proptest"))]
1915impl PropRelationDescDiff {
1916 pub fn apply(self, desc: &mut RelationDesc) {
1917 match self {
1918 PropRelationDescDiff::AddColumn { name, typ } => {
1919 let new_idx = desc.metadata.len();
1920 let meta = ColumnMetadata {
1921 name,
1922 typ_idx: new_idx,
1923 added: RelationVersion(0),
1924 dropped: None,
1925 };
1926 let prev = desc.metadata.insert(ColumnIndex(new_idx), meta);
1927 desc.typ.column_types.push(typ);
1928
1929 assert_none!(prev);
1930 assert_eq!(desc.metadata.len(), desc.typ.column_types.len());
1931 }
1932 PropRelationDescDiff::DropColumn { name } => {
1933 let next_version = desc
1934 .metadata
1935 .values()
1936 .map(|meta| meta.dropped.unwrap_or(meta.added))
1937 .max()
1938 .unwrap_or_else(RelationVersion::root)
1939 .bump();
1940 let Some(metadata) = desc.metadata.values_mut().find(|meta| meta.name == name)
1941 else {
1942 return;
1943 };
1944 if metadata.dropped.is_none() {
1945 metadata.dropped = Some(next_version);
1946 }
1947 }
1948 PropRelationDescDiff::ToggleNullability { name } => {
1949 let Some((pos, _)) = desc.get_by_name(&name) else {
1950 return;
1951 };
1952 let col_type = desc
1953 .typ
1954 .column_types
1955 .get_mut(pos)
1956 .expect("ColumnNames and SqlColumnTypes out of sync!");
1957 col_type.nullable = !col_type.nullable;
1958 }
1959 PropRelationDescDiff::ChangeType { name, typ } => {
1960 let Some((pos, _)) = desc.get_by_name(&name) else {
1961 return;
1962 };
1963 let col_type = desc
1964 .typ
1965 .column_types
1966 .get_mut(pos)
1967 .expect("ColumnNames and SqlColumnTypes out of sync!");
1968 *col_type = typ;
1969 }
1970 }
1971 }
1972}
1973
1974#[cfg(any(test, feature = "proptest"))]
1976pub fn arb_relation_desc_diff(
1977 source: &RelationDesc,
1978) -> impl Strategy<Value = Vec<PropRelationDescDiff>> + use<> {
1979 let source = Rc::new(source.clone());
1980 let num_source_columns = source.typ.columns().len();
1981
1982 let num_add_columns = Union::new_weighted(vec![(100, Just(0..8)), (1, Just(8..64))]);
1983 let add_columns_strat = num_add_columns
1984 .prop_flat_map(|num_columns| {
1985 proptest::collection::vec((any::<ColumnName>(), any::<SqlColumnType>()), num_columns)
1986 })
1987 .prop_map(|cols| {
1988 cols.into_iter()
1989 .map(|(name, typ)| PropRelationDescDiff::AddColumn { name, typ })
1990 .collect::<Vec<_>>()
1991 });
1992
1993 if num_source_columns == 0 {
1995 return add_columns_strat.boxed();
1996 }
1997
1998 let source_ = Rc::clone(&source);
1999 let drop_columns_strat = (0..num_source_columns).prop_perturb(move |num_columns, mut rng| {
2000 let mut set = BTreeSet::default();
2001 for _ in 0..num_columns {
2002 let col_idx = rng.random_range(0..num_source_columns);
2003 set.insert(source_.get_name(col_idx).clone());
2004 }
2005 set.into_iter()
2006 .map(|name| PropRelationDescDiff::DropColumn { name })
2007 .collect::<Vec<_>>()
2008 });
2009
2010 let source_ = Rc::clone(&source);
2011 let toggle_nullability_strat =
2012 (0..num_source_columns).prop_perturb(move |num_columns, mut rng| {
2013 let mut set = BTreeSet::default();
2014 for _ in 0..num_columns {
2015 let col_idx = rng.random_range(0..num_source_columns);
2016 set.insert(source_.get_name(col_idx).clone());
2017 }
2018 set.into_iter()
2019 .map(|name| PropRelationDescDiff::ToggleNullability { name })
2020 .collect::<Vec<_>>()
2021 });
2022
2023 let source_ = Rc::clone(&source);
2024 let change_type_strat = (0..num_source_columns)
2025 .prop_perturb(move |num_columns, mut rng| {
2026 let mut set = BTreeSet::default();
2027 for _ in 0..num_columns {
2028 let col_idx = rng.random_range(0..num_source_columns);
2029 set.insert(source_.get_name(col_idx).clone());
2030 }
2031 set
2032 })
2033 .prop_flat_map(|cols| {
2034 proptest::collection::vec(any::<SqlColumnType>(), cols.len())
2035 .prop_map(move |types| (cols.clone(), types))
2036 })
2037 .prop_map(|(cols, types)| {
2038 cols.into_iter()
2039 .zip_eq(types)
2040 .map(|(name, typ)| PropRelationDescDiff::ChangeType { name, typ })
2041 .collect::<Vec<_>>()
2042 });
2043
2044 (
2045 add_columns_strat,
2046 drop_columns_strat,
2047 toggle_nullability_strat,
2048 change_type_strat,
2049 )
2050 .prop_map(|(adds, drops, toggles, changes)| {
2051 adds.into_iter()
2052 .chain(drops)
2053 .chain(toggles)
2054 .chain(changes)
2055 .collect::<Vec<_>>()
2056 })
2057 .prop_shuffle()
2058 .boxed()
2059}
2060
2061#[cfg(test)]
2062mod tests {
2063 use super::*;
2064 use prost::Message;
2065
2066 #[mz_ore::test]
2070 #[should_panic(expected = "dense RelationDesc")]
2071 fn apply_demand_rejects_non_dense_desc() {
2072 let desc = RelationDesc::builder()
2073 .with_column("a", SqlScalarType::Int32.nullable(false))
2074 .with_column("b", SqlScalarType::Int32.nullable(false))
2075 .with_column("c", SqlScalarType::Int32.nullable(false))
2076 .finish();
2077 let mut versioned = VersionedRelationDesc::new(desc);
2078 let version = versioned.drop_column("b");
2079 let desc = versioned.at_version(RelationVersionSelector::Specific(version));
2080 let _ = desc.apply_demand(&BTreeSet::from([0]));
2081 }
2082
2083 #[mz_ore::test]
2084 #[cfg_attr(miri, ignore)] fn smoktest_at_version() {
2086 let desc = RelationDesc::builder()
2087 .with_column("a", SqlScalarType::Bool.nullable(true))
2088 .with_column("z", SqlScalarType::String.nullable(false))
2089 .finish();
2090
2091 let mut versioned_desc = VersionedRelationDesc {
2092 inner: desc.clone(),
2093 };
2094 versioned_desc.validate();
2095
2096 let latest = versioned_desc.at_version(RelationVersionSelector::Latest);
2097 assert_eq!(desc, latest);
2098
2099 let v0 = versioned_desc.at_version(RelationVersionSelector::specific(0));
2100 assert_eq!(desc, v0);
2101
2102 let v3 = versioned_desc.at_version(RelationVersionSelector::specific(3));
2103 assert_eq!(desc, v3);
2104
2105 let v1 = versioned_desc.add_column("b", SqlScalarType::Bytes.nullable(false));
2106 assert_eq!(v1, RelationVersion(1));
2107
2108 let v1 = versioned_desc.at_version(RelationVersionSelector::Specific(v1));
2109 insta::assert_json_snapshot!(v1.metadata, @r###"
2110 {
2111 "0": {
2112 "name": "a",
2113 "typ_idx": 0,
2114 "added": 0,
2115 "dropped": null
2116 },
2117 "1": {
2118 "name": "z",
2119 "typ_idx": 1,
2120 "added": 0,
2121 "dropped": null
2122 },
2123 "2": {
2124 "name": "b",
2125 "typ_idx": 2,
2126 "added": 1,
2127 "dropped": null
2128 }
2129 }
2130 "###);
2131
2132 let v0_b = versioned_desc.at_version(RelationVersionSelector::specific(0));
2134 assert!(v0.iter().eq(v0_b.iter()));
2135
2136 let v2 = versioned_desc.drop_column("z");
2137 assert_eq!(v2, RelationVersion(2));
2138
2139 let v2 = versioned_desc.at_version(RelationVersionSelector::Specific(v2));
2140 insta::assert_json_snapshot!(v2.metadata, @r###"
2141 {
2142 "0": {
2143 "name": "a",
2144 "typ_idx": 0,
2145 "added": 0,
2146 "dropped": null
2147 },
2148 "2": {
2149 "name": "b",
2150 "typ_idx": 1,
2151 "added": 1,
2152 "dropped": null
2153 }
2154 }
2155 "###);
2156
2157 let v0_c = versioned_desc.at_version(RelationVersionSelector::specific(0));
2159 assert!(v0.iter().eq(v0_c.iter()));
2160
2161 let v1_b = versioned_desc.at_version(RelationVersionSelector::specific(1));
2162 assert!(v1.iter().eq(v1_b.iter()));
2163
2164 insta::assert_json_snapshot!(versioned_desc.inner.metadata, @r###"
2165 {
2166 "0": {
2167 "name": "a",
2168 "typ_idx": 0,
2169 "added": 0,
2170 "dropped": null
2171 },
2172 "1": {
2173 "name": "z",
2174 "typ_idx": 1,
2175 "added": 0,
2176 "dropped": 2
2177 },
2178 "2": {
2179 "name": "b",
2180 "typ_idx": 2,
2181 "added": 1,
2182 "dropped": null
2183 }
2184 }
2185 "###);
2186 }
2187
2188 #[mz_ore::test]
2189 #[cfg_attr(miri, ignore)] fn test_dropping_columns_with_keys() {
2191 let desc = RelationDesc::builder()
2192 .with_column("a", SqlScalarType::Bool.nullable(true))
2193 .with_column("z", SqlScalarType::String.nullable(false))
2194 .with_key(vec![1])
2195 .finish();
2196
2197 let mut versioned_desc = VersionedRelationDesc {
2198 inner: desc.clone(),
2199 };
2200 versioned_desc.validate();
2201
2202 let v1 = versioned_desc.drop_column("a");
2203 assert_eq!(v1, RelationVersion(1));
2204
2205 let v1 = versioned_desc.at_version(RelationVersionSelector::Specific(v1));
2207 insta::assert_json_snapshot!(v1, @r###"
2208 {
2209 "typ": {
2210 "column_types": [
2211 {
2212 "scalar_type": "String",
2213 "nullable": false
2214 }
2215 ],
2216 "keys": [
2217 [
2218 0
2219 ]
2220 ]
2221 },
2222 "metadata": {
2223 "1": {
2224 "name": "z",
2225 "typ_idx": 0,
2226 "added": 0,
2227 "dropped": null
2228 }
2229 }
2230 }
2231 "###);
2232
2233 let v0 = versioned_desc.at_version(RelationVersionSelector::specific(0));
2235 insta::assert_json_snapshot!(v0, @r###"
2236 {
2237 "typ": {
2238 "column_types": [
2239 {
2240 "scalar_type": "Bool",
2241 "nullable": true
2242 },
2243 {
2244 "scalar_type": "String",
2245 "nullable": false
2246 }
2247 ],
2248 "keys": [
2249 [
2250 1
2251 ]
2252 ]
2253 },
2254 "metadata": {
2255 "0": {
2256 "name": "a",
2257 "typ_idx": 0,
2258 "added": 0,
2259 "dropped": 1
2260 },
2261 "1": {
2262 "name": "z",
2263 "typ_idx": 1,
2264 "added": 0,
2265 "dropped": null
2266 }
2267 }
2268 }
2269 "###);
2270 }
2271
2272 #[mz_ore::test]
2273 #[cfg_attr(miri, ignore)] fn roundtrip_relation_desc_without_metadata() {
2275 let typ = ProtoRelationType {
2276 column_types: vec![
2277 SqlScalarType::String.nullable(false).into_proto(),
2278 SqlScalarType::Bool.nullable(true).into_proto(),
2279 ],
2280 keys: vec![],
2281 };
2282 let proto = ProtoRelationDesc {
2283 typ: Some(typ),
2284 names: vec![
2285 ColumnName("a".into()).into_proto(),
2286 ColumnName("b".into()).into_proto(),
2287 ],
2288 metadata: vec![],
2289 };
2290 let desc: RelationDesc = proto.into_rust().unwrap();
2291
2292 insta::assert_json_snapshot!(desc, @r###"
2293 {
2294 "typ": {
2295 "column_types": [
2296 {
2297 "scalar_type": "String",
2298 "nullable": false
2299 },
2300 {
2301 "scalar_type": "Bool",
2302 "nullable": true
2303 }
2304 ],
2305 "keys": []
2306 },
2307 "metadata": {
2308 "0": {
2309 "name": "a",
2310 "typ_idx": 0,
2311 "added": 0,
2312 "dropped": null
2313 },
2314 "1": {
2315 "name": "b",
2316 "typ_idx": 1,
2317 "added": 0,
2318 "dropped": null
2319 }
2320 }
2321 }
2322 "###);
2323 }
2324
2325 #[mz_ore::test]
2326 #[should_panic(expected = "column named 'a' already exists!")]
2327 fn test_add_column_with_same_name_panics() {
2328 let desc = RelationDesc::builder()
2329 .with_column("a", SqlScalarType::Bool.nullable(true))
2330 .finish();
2331 let mut versioned = VersionedRelationDesc::new(desc);
2332
2333 let _ = versioned.add_column("a", SqlScalarType::String.nullable(false));
2334 }
2335
2336 #[mz_ore::test]
2337 #[cfg_attr(miri, ignore)] fn test_add_column_with_same_name_prev_dropped() {
2339 let desc = RelationDesc::builder()
2340 .with_column("a", SqlScalarType::Bool.nullable(true))
2341 .finish();
2342 let mut versioned = VersionedRelationDesc::new(desc);
2343
2344 let v1 = versioned.drop_column("a");
2345 let v1 = versioned.at_version(RelationVersionSelector::Specific(v1));
2346 insta::assert_json_snapshot!(v1, @r###"
2347 {
2348 "typ": {
2349 "column_types": [],
2350 "keys": []
2351 },
2352 "metadata": {}
2353 }
2354 "###);
2355
2356 let v2 = versioned.add_column("a", SqlScalarType::String.nullable(false));
2357 let v2 = versioned.at_version(RelationVersionSelector::Specific(v2));
2358 insta::assert_json_snapshot!(v2, @r###"
2359 {
2360 "typ": {
2361 "column_types": [
2362 {
2363 "scalar_type": "String",
2364 "nullable": false
2365 }
2366 ],
2367 "keys": []
2368 },
2369 "metadata": {
2370 "1": {
2371 "name": "a",
2372 "typ_idx": 0,
2373 "added": 2,
2374 "dropped": null
2375 }
2376 }
2377 }
2378 "###);
2379 }
2380
2381 #[mz_ore::test]
2382 #[cfg_attr(miri, ignore)] fn test_drop_column_added_after_root() {
2384 let desc = RelationDesc::builder()
2385 .with_column("a", SqlScalarType::Bool.nullable(true))
2386 .finish();
2387 let mut versioned = VersionedRelationDesc::new(desc);
2388
2389 let v1 = versioned.add_column("b", SqlScalarType::String.nullable(false));
2390 let v2 = versioned.drop_column("b");
2391 assert_eq!(v1, RelationVersion(1));
2392 assert_eq!(v2, RelationVersion(2));
2393
2394 assert_eq!(
2395 versioned
2396 .at_version(RelationVersionSelector::Specific(v1))
2397 .arity(),
2398 2
2399 );
2400 assert_eq!(
2401 versioned
2402 .at_version(RelationVersionSelector::Specific(v2))
2403 .arity(),
2404 1
2405 );
2406 }
2407
2408 #[mz_ore::test]
2409 #[cfg_attr(miri, ignore)] fn relation_desc_proto_rejects_corrupt_shapes() {
2411 fn proto(num_types: usize, num_names: usize, keys: Vec<Vec<usize>>) -> ProtoRelationDesc {
2412 let mut typ = SqlRelationType::new(vec![SqlScalarType::Bool.nullable(true); num_types]);
2413 typ.keys = keys;
2414 ProtoRelationDesc {
2415 typ: Some(typ.into_proto()),
2416 names: (0..num_names)
2417 .map(|i| ColumnName::from(format!("c{i}")).into_proto())
2418 .collect(),
2419 metadata: vec![],
2420 }
2421 }
2422
2423 for (num_types, num_names) in [(0, 1), (2, 0), (3, 1)] {
2427 let err = RelationDesc::from_proto(proto(num_types, num_names, vec![]))
2428 .expect_err("length mismatch must be rejected");
2429 assert!(err.to_string().contains("length mismatch"), "{err}");
2430 }
2431 let err = RelationDesc::from_proto(proto(1, 1, vec![vec![7]]))
2432 .expect_err("out of bounds key must be rejected");
2433 assert!(err.to_string().contains("out of bounds"), "{err}");
2434
2435 let desc = RelationDesc::from_proto(proto(2, 2, vec![vec![1]])).expect("valid");
2437 assert_eq!(desc.iter().count(), 2);
2438 }
2439
2440 #[mz_ore::test]
2441 #[cfg_attr(miri, ignore)]
2442 fn apply_demand() {
2443 let desc = RelationDesc::builder()
2444 .with_column("a", SqlScalarType::String.nullable(true))
2445 .with_column("b", SqlScalarType::Int64.nullable(false))
2446 .with_column("c", SqlScalarType::Time.nullable(false))
2447 .finish();
2448 let desc = desc.apply_demand(&BTreeSet::from([0, 2]));
2449 assert_eq!(desc.arity(), 2);
2450 VersionedRelationDesc::new(desc).validate();
2452 }
2453
2454 #[mz_ore::test]
2455 #[cfg_attr(miri, ignore)]
2456 fn smoketest_column_index_stable_ident() {
2457 let idx_a = ColumnIndex(42);
2458 assert_eq!(idx_a.to_stable_name(), "42");
2460 }
2461
2462 #[mz_ore::test]
2463 #[cfg_attr(miri, ignore)] fn proptest_relation_desc_roundtrips() {
2465 fn testcase(og: RelationDesc) {
2466 let bytes = og.into_proto().encode_to_vec();
2467 let proto = ProtoRelationDesc::decode(&bytes[..]).unwrap();
2468 let rnd = RelationDesc::from_proto(proto).unwrap();
2469
2470 assert_eq!(og, rnd);
2471 }
2472
2473 proptest!(|(desc in any::<RelationDesc>())| {
2474 testcase(desc);
2475 });
2476
2477 let strat = any::<RelationDesc>().prop_flat_map(|desc| {
2478 arb_relation_desc_diff(&desc).prop_map(move |diffs| (desc.clone(), diffs))
2479 });
2480
2481 proptest!(|((mut desc, diffs) in strat)| {
2482 for diff in diffs {
2483 diff.apply(&mut desc);
2484 };
2485 testcase(desc);
2486 });
2487 }
2488}