Skip to content

Commit

Permalink
Sometimes we can give the type to MRE::take_safely
Browse files Browse the repository at this point in the history
  • Loading branch information
ggevay committed Dec 20, 2024
1 parent 233f333 commit 924528b
Show file tree
Hide file tree
Showing 4 changed files with 50 additions and 7 deletions.
44 changes: 42 additions & 2 deletions src/expr/src/relation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ use mz_ore::collections::CollectionExt;
use mz_ore::id_gen::IdGen;
use mz_ore::metrics::Histogram;
use mz_ore::num::NonNeg;
use mz_ore::soft_assert_eq_no_log;
use mz_ore::stack::RecursionLimitError;
use mz_ore::str::Indent;
use mz_proto::{IntoRustIfSome, ProtoType, RustType, TryFromProtoError};
Expand Down Expand Up @@ -1688,8 +1689,44 @@ impl MirRelationExpr {
}

/// Take ownership of `self`, leaving an empty `MirRelationExpr::Constant` with the correct type.
///
/// This calls `.typ()` to determine scalar types. If you already know the scalar types type,
/// then use either
/// [MirRelationExpr::take_safely_with_col_types] or [MirRelationExpr::take_safely_with_rel_type].
pub fn take_safely(&mut self) -> MirRelationExpr {
let typ = self.typ();
self.take_safely_with_rel_type(self.typ())
}

/// Take ownership of `self`, leaving an empty `MirRelationExpr::Constant` with the given scalar
/// types. Keys and nullability are ignored in the given `RelationType`, and instead we set the
/// best possible key and nullability, since we are making an empty collection.
///
/// Compared to `take_safely()`, this just saves the cost of the `.typ()` call.
pub fn take_safely_with_col_types(&mut self, typ: Vec<ColumnType>) -> MirRelationExpr {
self.take_safely_with_rel_type(RelationType::new(typ))
}

/// Take ownership of `self`, leaving an empty `MirRelationExpr::Constant` with the given scalar
/// types. Keys and nullability are ignored in the given `RelationType`, and instead we set the
/// best possible key and nullability, since we are making an empty collection.
///
/// Compared to `take_safely()`, this just saves the cost of the `.typ()` call.
pub fn take_safely_with_rel_type(&mut self, mut typ: RelationType) -> MirRelationExpr {
soft_assert_eq_no_log!(
self.typ()
.column_types
.iter()
.map(|ct| ct.scalar_type.clone())
.collect_vec(),
typ.column_types
.iter()
.map(|ct| ct.scalar_type.clone())
.collect_vec()
);
typ.keys = vec![vec![]];
for ct in typ.column_types.iter_mut() {
ct.nullable = false;
}
std::mem::replace(
self,
MirRelationExpr::Constant {
Expand All @@ -1698,6 +1735,7 @@ impl MirRelationExpr {
},
)
}

/// Take ownership of `self`, leaving an empty `MirRelationExpr::Constant` with an **incorrect** type.
///
/// This should only be used if `self` is about to be dropped or otherwise overwritten.
Expand Down Expand Up @@ -2233,11 +2271,13 @@ impl MirRelationExpr {
value.visit_pre_mut(|e| {
if let MirRelationExpr::Get {
id: crate::Id::Local(id),
typ,
..
} = e
{
let typ = typ.clone();
if deadlist.contains(id) {
e.take_safely();
e.take_safely_with_rel_type(typ);
}
}
});
Expand Down
5 changes: 3 additions & 2 deletions src/transform/src/equivalence_propagation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,7 @@ impl EquivalencePropagation {
let expr_type = derived
.value::<RelationType>()
.expect("RelationType required");
assert!(expr_type.is_some());
let expr_equivalences = derived
.value::<Equivalences>()
.expect("Equivalences required");
Expand All @@ -132,7 +133,7 @@ impl EquivalencePropagation {
let expr_equivalences = if let Some(e) = expr_equivalences {
e
} else {
expr.take_safely();
expr.take_safely_with_col_types(expr_type.clone().unwrap());
return;
};

Expand All @@ -147,7 +148,7 @@ impl EquivalencePropagation {

outer_equivalences.minimize(expr_type.as_ref().map(|x| &x[..]));
if outer_equivalences.unsatisfiable() {
expr.take_safely();
expr.take_safely_with_col_types(expr_type.clone().unwrap());
return;
}

Expand Down
4 changes: 2 additions & 2 deletions src/transform/src/fold_constants.rs
Original file line number Diff line number Diff line change
Expand Up @@ -250,7 +250,7 @@ impl FoldConstants {
.iter()
.any(|p| p.is_literal_false() || p.is_literal_null())
{
relation.take_safely();
relation.take_safely_with_rel_type(relation_type.clone());
} else if let Some((rows, ..)) = (**input).as_const() {
// Evaluate errors last, to reduce risk of spurious errors.
predicates.sort_by_key(|p| p.is_literal_err());
Expand Down Expand Up @@ -291,7 +291,7 @@ impl FoldConstants {
..
} => {
if inputs.iter().any(|e| e.is_empty()) {
relation.take_safely();
relation.take_safely_with_rel_type(relation_type.clone());
} else if let Some(e) = inputs.iter().find_map(|i| i.as_const_err()) {
*relation = MirRelationExpr::Constant {
rows: Err(e.clone()),
Expand Down
4 changes: 3 additions & 1 deletion src/transform/src/predicate_pushdown.rs
Original file line number Diff line number Diff line change
Expand Up @@ -595,7 +595,9 @@ impl PredicatePushdown {
.count()
> 1
{
relation.take_safely();
relation.take_safely_with_rel_type(
relation.typ_with_input_types(&input_types),
);
return Ok(());
}

Expand Down

0 comments on commit 924528b

Please sign in to comment.