1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317
// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0.
//! Planning of linear joins.
use std::collections::BTreeMap;
use mz_expr::{
join_permutations, permutation_for_arrangement, JoinInputCharacteristics, MapFilterProject,
MirScalarExpr,
};
use mz_proto::{IntoRustIfSome, ProtoType, RustType, TryFromProtoError};
use proptest::prelude::*;
use proptest::result::Probability;
use proptest_derive::Arbitrary;
use serde::{Deserialize, Serialize};
use crate::plan::join::{
JoinBuildState, JoinClosure, ProtoLinearJoinPlan, ProtoLinearStagePlan, ProtoMirScalarVec,
};
use crate::plan::AvailableCollections;
/// A plan for the execution of a linear join.
///
/// A linear join is a sequence of stages, each of which introduces
/// a new collection. Each stage is represented by a [LinearStagePlan].
#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq, Ord, PartialOrd)]
pub struct LinearJoinPlan {
/// The source relation from which we start the join.
pub source_relation: usize,
/// The arrangement to use for the source relation, if any
pub source_key: Option<Vec<MirScalarExpr>>,
/// An initial closure to apply before any stages.
///
/// Values of `None` indicate the identity closure.
pub initial_closure: Option<JoinClosure>,
/// A *sequence* of stages to apply one after the other.
pub stage_plans: Vec<LinearStagePlan>,
/// A concluding closure to apply after the last stage.
///
/// Values of `None` indicate the identity closure.
pub final_closure: Option<JoinClosure>,
}
impl Arbitrary for LinearJoinPlan {
type Parameters = ();
type Strategy = BoxedStrategy<Self>;
fn arbitrary_with(_: Self::Parameters) -> Self::Strategy {
(
any::<usize>(),
any_with::<Option<Vec<MirScalarExpr>>>((Probability::default(), ((0..3).into(), ()))),
any::<Option<JoinClosure>>(),
prop::collection::vec(any::<LinearStagePlan>(), 0..3),
any::<Option<JoinClosure>>(),
)
.prop_map(
|(source_relation, source_key, initial_closure, stage_plans, final_closure)| {
LinearJoinPlan {
source_relation,
source_key,
initial_closure,
stage_plans,
final_closure,
}
},
)
.boxed()
}
}
impl RustType<ProtoLinearJoinPlan> for LinearJoinPlan {
fn into_proto(&self) -> ProtoLinearJoinPlan {
ProtoLinearJoinPlan {
source_relation: self.source_relation.into_proto(),
source_key: self.source_key.into_proto(),
initial_closure: self.initial_closure.into_proto(),
stage_plans: self.stage_plans.into_proto(),
final_closure: self.final_closure.into_proto(),
}
}
fn from_proto(proto: ProtoLinearJoinPlan) -> Result<Self, TryFromProtoError> {
Ok(LinearJoinPlan {
source_relation: proto.source_relation.into_rust()?,
source_key: proto.source_key.into_rust()?,
initial_closure: proto.initial_closure.into_rust()?,
stage_plans: proto.stage_plans.into_rust()?,
final_closure: proto.final_closure.into_rust()?,
})
}
}
impl RustType<ProtoMirScalarVec> for Vec<MirScalarExpr> {
fn into_proto(&self) -> ProtoMirScalarVec {
ProtoMirScalarVec {
values: self.into_proto(),
}
}
fn from_proto(proto: ProtoMirScalarVec) -> Result<Self, TryFromProtoError> {
proto.values.into_rust()
}
}
/// A plan for the execution of one stage of a linear join.
///
/// Each stage is a binary join between the current accumulated
/// join results, and a new collection. The former is referred to
/// as the "stream" and the latter the "lookup".
#[derive(Arbitrary, Clone, Debug, Serialize, Deserialize, Eq, PartialEq, Ord, PartialOrd)]
pub struct LinearStagePlan {
/// The index of the relation into which we will look up.
pub lookup_relation: usize,
/// The key expressions to use for the stream relation.
pub stream_key: Vec<MirScalarExpr>,
/// Columns to retain from the stream relation.
/// These columns are those that are not redundant with `stream_key`,
/// and cannot be read out of the key component of an arrangement.
pub stream_thinning: Vec<usize>,
/// The key expressions to use for the lookup relation.
pub lookup_key: Vec<MirScalarExpr>,
/// The closure to apply to the concatenation of the key columns,
/// the stream value columns, and the lookup value colunms.
pub closure: JoinClosure,
}
impl RustType<ProtoLinearStagePlan> for LinearStagePlan {
fn into_proto(&self) -> ProtoLinearStagePlan {
ProtoLinearStagePlan {
lookup_relation: self.lookup_relation.into_proto(),
stream_key: self.stream_key.into_proto(),
stream_thinning: self.stream_thinning.into_proto(),
lookup_key: self.lookup_key.into_proto(),
closure: Some(self.closure.into_proto()),
}
}
fn from_proto(proto: ProtoLinearStagePlan) -> Result<Self, TryFromProtoError> {
Ok(Self {
lookup_relation: proto.lookup_relation.into_rust()?,
stream_key: proto.stream_key.into_rust()?,
stream_thinning: proto.stream_thinning.into_rust()?,
lookup_key: proto.lookup_key.into_rust()?,
closure: proto
.closure
.into_rust_if_some("ProtoLinearStagePlan::closure")?,
})
}
}
impl LinearJoinPlan {
/// Create a new join plan from the required arguments.
pub fn create_from(
source_relation: usize,
// When specified, a key and its corresponding permutation and thinning.
source_arrangement: Option<&(Vec<MirScalarExpr>, BTreeMap<usize, usize>, Vec<usize>)>,
equivalences: &[Vec<MirScalarExpr>],
join_order: &[(usize, Vec<MirScalarExpr>, Option<JoinInputCharacteristics>)],
input_mapper: mz_expr::JoinInputMapper,
// An MFP to apply to the result of the join.
mfp_above: &mut MapFilterProject,
available: &[AvailableCollections],
) -> (Self, Vec<AvailableCollections>) {
let mut requested: Vec<AvailableCollections> =
vec![Default::default(); input_mapper.total_inputs()];
// Temporal predicates cannot currently be pushed down, and so we extract them and
// set `mfp` to be the temporal predicates at the end of the method.
let temporal_mfp = mfp_above.extract_temporal();
// Construct initial join build state.
// This state evolves as we build the join dataflow.
let mut join_build_state = JoinBuildState::new(
input_mapper.global_columns(source_relation),
equivalences,
mfp_above,
);
// Determine if `available` contains the indicated source arrangement, and if not request it.
if let Some(kpt) = source_arrangement {
if !available[source_relation].arranged.contains(kpt) {
requested[source_relation].arranged.push(kpt.clone());
}
} else if !available[source_relation].raw {
requested[source_relation].raw = true;
}
// Tracks the current state of a potentially arranged stream input.
// Currently only possible for the first input, if it starts arranged.
let mut stream_kpt = source_arrangement.cloned();
let mut unthinned_stream_arity = input_mapper.input_arity(source_relation);
// Sequence of steps to apply.
let mut stage_plans = Vec::with_capacity(join_order.len());
// Track the set of bound input relations, for equivalence resolution.
let mut bound_inputs = vec![source_relation];
// Iterate through the join order instructions, assembling keys and
// closures to use.
for (lookup_relation, lookup_key, _characteristics) in join_order.iter() {
let available = &available[*lookup_relation];
let (lookup_permutation, lookup_thinning) = available
.arranged
.iter()
.find_map(|(key, permutation, thinning)| {
if key == lookup_key {
Some((permutation.clone(), thinning.clone()))
} else {
None
}
})
.unwrap_or_else(|| {
let (permutation, thinning) = permutation_for_arrangement(
lookup_key,
input_mapper.input_arity(*lookup_relation),
);
requested[*lookup_relation].arranged.push((
lookup_key.clone(),
permutation.clone(),
thinning.clone(),
));
(permutation, thinning)
});
// rebase the intended key to use global column identifiers.
let lookup_key_rebased = lookup_key
.iter()
.map(|k| input_mapper.map_expr_to_global(k.clone(), *lookup_relation))
.collect::<Vec<_>>();
// If the stream input is already arranged, we have key, permutation, and thinning already specified.
let (stream_key, stream_permutation, stream_thinning) =
stream_kpt.take().unwrap_or_else(|| {
// Expressions to use as a key for the stream of incoming updates
// are determined by locating the elements of `lookup_key` among
// the existing bound `columns`. If that cannot be done, the plan
// is irrecoverably defective and we panic.
// TODO: explicitly validate this before rendering.
let stream_key = lookup_key_rebased
.iter()
.map(|expr| {
let mut bound_expr = input_mapper
.find_bound_expr(
expr,
&bound_inputs,
&join_build_state.equivalences,
)
.expect("Expression in join plan is not bound at time of use");
// Rewrite column references to physical locations.
bound_expr.permute_map(&join_build_state.column_map);
bound_expr
})
.collect::<Vec<_>>();
let (stream_permutation, stream_thinning) =
permutation_for_arrangement(&stream_key, unthinned_stream_arity);
(stream_key, stream_permutation, stream_thinning)
});
let key_arity = stream_key.len();
let permutation = join_permutations(
key_arity,
stream_permutation,
stream_thinning.len(),
lookup_permutation,
);
// Introduce new columns and expressions they enable. Form a new closure.
let closure = join_build_state.add_columns(
input_mapper.global_columns(*lookup_relation),
&lookup_key_rebased,
key_arity + stream_thinning.len() + lookup_thinning.len(),
permutation,
);
let new_unthinned_stream_arity = closure.before.projection.len();
bound_inputs.push(*lookup_relation);
// record the stage plan as next in the path.
stage_plans.push(LinearStagePlan {
lookup_relation: *lookup_relation,
stream_key,
stream_thinning,
lookup_key: lookup_key.to_vec(),
closure,
});
unthinned_stream_arity = new_unthinned_stream_arity;
}
// determine a final closure, and complete the path plan.
let final_closure = join_build_state.complete();
let final_closure = if final_closure.is_identity() {
None
} else {
Some(final_closure)
};
// Now that `map_filter_project` has been captured in the state builder,
// assign the remaining temporal predicates to it, for the caller's use.
*mfp_above = temporal_mfp;
// Form and return the complete join plan.
let plan = LinearJoinPlan {
source_relation,
source_key: source_arrangement.map(|(key, _, _)| key.clone()),
initial_closure: None,
stage_plans,
final_closure,
};
(plan, requested)
}
}