use rudb_common::Result;
use rudb_plan::{BuildSide, ExprRef, JoinKind, Node, NodeRef, Plan};
use crate::estimate::{self, Statistics};
use crate::pass::{Context, Pass};
use crate::tables::{TableSet, Tables, produced};
use crate::walk;
#[derive(Debug, Clone, Copy)]
pub struct JoinOrder;
impl Pass for JoinOrder {
fn name(&self) -> &'static str {
"join_order"
}
fn run(&self, plan: &mut Plan, context: &Context) -> Result<()> {
reorder(plan, context.statistics());
Ok(())
}
}
pub fn reorder(plan: &mut Plan, stats: &Statistics) {
let mut tables = Tables::new();
let root = rebuild(plan, plan.root(), &mut tables, stats);
plan.set_root(root);
}
struct Part {
build: usize,
tables: TableSet,
rows: u64,
}
enum Build {
Leaf(NodeRef),
Pair { left: usize, right: usize, conditions: Vec<ExprRef> },
}
fn rebuild(plan: &mut Plan, at: NodeRef, tables: &mut Tables, stats: &Statistics) -> NodeRef {
if joining(plan, at) {
let mut leaves = Vec::new();
let mut conditions = Vec::new();
gather(plan, at, &mut leaves, &mut conditions);
let rebuilt: Vec<NodeRef> =
leaves.iter().map(|&leaf| rebuild(plan, leaf, tables, stats)).collect();
let chosen = match leaves.len() {
0..=2 => None,
_ => order(plan, at, &rebuilt, &conditions, tables, stats),
};
if let Some(chosen) = chosen {
return chosen;
}
if rebuilt == leaves {
return at;
}
return restack(plan, at, &leaves, &rebuilt);
}
let children: Vec<NodeRef> = plan.node(at).children().into_iter().flatten().collect();
let rebuilt: Vec<NodeRef> =
children.iter().map(|&child| rebuild(plan, child, tables, stats)).collect();
if rebuilt == children {
return at;
}
let mut node = plan.node(at).clone();
walk::replace_children(&mut node, &rebuilt);
plan.add_node(node)
}
fn joining(plan: &Plan, at: NodeRef) -> bool {
matches!(*plan.node(at), Node::CrossProduct { .. } | Node::Join { kind: JoinKind::Inner, .. })
}
fn gather(plan: &Plan, at: NodeRef, leaves: &mut Vec<NodeRef>, conditions: &mut Vec<ExprRef>) {
match *plan.node(at) {
Node::CrossProduct { left, right } => {
gather(plan, left, leaves, conditions);
gather(plan, right, leaves, conditions);
}
Node::Join { left, right, kind: JoinKind::Inner, conditions: list, .. } => {
conditions.extend_from_slice(plan.expr_list(list));
gather(plan, left, leaves, conditions);
gather(plan, right, leaves, conditions);
}
_ => leaves.push(at),
}
}
fn restack(plan: &mut Plan, at: NodeRef, leaves: &[NodeRef], rebuilt: &[NodeRef]) -> NodeRef {
if let Some(found) = leaves.iter().position(|&leaf| leaf == at) {
return rebuilt[found];
}
let children: Vec<NodeRef> = plan.node(at).children().into_iter().flatten().collect();
let children: Vec<NodeRef> =
children.into_iter().map(|child| restack(plan, child, leaves, rebuilt)).collect();
let mut node = plan.node(at).clone();
walk::replace_children(&mut node, &children);
plan.add_node(node)
}
fn order(
plan: &mut Plan,
at: NodeRef,
leaves: &[NodeRef],
conditions: &[ExprRef],
tables: &mut Tables,
stats: &Statistics,
) -> Option<NodeRef> {
let mut builds: Vec<Build> = leaves.iter().map(|&leaf| Build::Leaf(leaf)).collect();
let mut parts = Vec::with_capacity(leaves.len());
for (build, &leaf) in leaves.iter().enumerate() {
parts.push(Part {
build,
tables: produced(plan, leaf),
rows: estimate::rows(plan, leaf, stats)?,
});
}
let mut whole = TableSet::new();
for part in &parts {
whole.extend(&part.tables);
}
let mut pending: Vec<(ExprRef, TableSet)> = Vec::with_capacity(conditions.len());
for &condition in conditions {
let reads = tables.of(plan, condition);
if !reads.is_subset_of(&whole) || parts.iter().any(|part| reads.is_subset_of(&part.tables))
{
return None;
}
pending.push((condition, reads));
}
let (_, before, was) = cost(plan, at, stats)?;
let mut after = 0u64;
let mut built = 0usize;
while parts.len() > 1 {
let (left, right, rows) = cheapest(&parts, &pending);
let mut union = parts[left].tables.clone();
union.extend(&parts[right].tables);
let conditions: Vec<ExprRef> = pending
.iter()
.filter(|(_, reads)| reads.is_subset_of(&union))
.map(|(condition, _)| *condition)
.collect();
pending.retain(|(_, reads)| !reads.is_subset_of(&union));
let right = parts.remove(right);
let left = parts.remove(left);
built += usize::from(conditions.is_empty());
builds.push(Build::Pair { left: left.build, right: right.build, conditions });
after = after.saturating_add(rows);
parts.push(Part { build: builds.len() - 1, tables: union, rows });
}
if built >= was || after >= before {
return None;
}
Some(put(plan, &builds, parts[0].build))
}
fn put(plan: &mut Plan, builds: &[Build], at: usize) -> NodeRef {
match &builds[at] {
Build::Leaf(node) => *node,
Build::Pair { left, right, conditions } => {
let conditions = conditions.clone();
let left = put(plan, builds, *left);
let right = put(plan, builds, *right);
if conditions.is_empty() {
return plan.add_node(Node::CrossProduct { left, right });
}
let conditions = plan.add_expr_list(&conditions);
plan.add_node(Node::Join {
left,
right,
kind: JoinKind::Inner,
conditions,
build: BuildSide::default(),
})
}
}
}
fn cheapest(parts: &[Part], pending: &[(ExprRef, TableSet)]) -> (usize, usize, u64) {
let mut best: Option<Pick> = None;
for left in 0..parts.len() {
for right in left + 1..parts.len() {
let mut union = parts[left].tables.clone();
union.extend(&parts[right].tables);
let linked = pending.iter().any(|(_, reads)| reads.is_subset_of(&union));
let rows = if linked {
parts[left].rows.max(parts[right].rows)
} else {
parts[left].rows.saturating_mul(parts[right].rows)
};
let order = (!linked, rows, parts[left].rows.saturating_add(parts[right].rows));
if best.is_none_or(|held| order < held.order) {
best = Some(Pick { left, right, rows, order });
}
}
}
let best = best.expect("a region has at least two parts");
(best.left, best.right, best.rows)
}
#[derive(Clone, Copy)]
struct Pick {
left: usize,
right: usize,
rows: u64,
order: (bool, u64, u64),
}
fn cost(plan: &Plan, at: NodeRef, stats: &Statistics) -> Option<(u64, u64, usize)> {
let (left, right, linked) = match *plan.node(at) {
Node::CrossProduct { left, right } => (left, right, false),
Node::Join { left, right, kind: JoinKind::Inner, conditions, .. } => {
(left, right, !plan.expr_list(conditions).is_empty())
}
_ => return Some((estimate::rows(plan, at, stats)?, 0, 0)),
};
let (left, under_left, crossed_left) = cost(plan, left, stats)?;
let (right, under_right, crossed_right) = cost(plan, right, stats)?;
let rows = if linked { left.max(right) } else { left.saturating_mul(right) };
Some((
rows,
under_left.saturating_add(under_right).saturating_add(rows),
crossed_left + crossed_right + usize::from(!linked),
))
}
#[cfg(test)]
mod tests {
use rudb_plan::Plan;
use crate::estimate::Statistics;
use super::reorder;
fn ordered(text: &str) -> String {
let mut counts = Statistics::new();
for (table, rows) in [("t", 1000), ("u", 10), ("v", 100), ("w", 100_000)] {
counts.record("memory", "main", table, rows);
}
let mut plan =
Plan::parse(text).unwrap_or_else(|error| panic!("{text} did not parse: {error}"));
reorder(&mut plan, &counts);
plan.validate().unwrap_or_else(|error| panic!("{text} did not stay valid: {error}"));
plan.to_string()
}
#[test]
fn a_cross_product_the_conditions_can_avoid_is_not_built() {
assert_eq!(
ordered(concat!(
"Join INNER on=[(#1.0::BIGINT = #0.0::BIGINT)::BOOLEAN, ",
"(#1.0::BIGINT = #2.0::BIGINT)::BOOLEAN]\n",
" CrossProduct\n",
" Get memory.main.t AS t #0 [a::BIGINT]\n",
" Get memory.main.v AS v #2 [c::BIGINT]\n",
" Get memory.main.u AS u #1 [b::BIGINT]\n",
)),
concat!(
"Join INNER on=[(#1.0::BIGINT = #0.0::BIGINT)::BOOLEAN]\n",
" Get memory.main.t AS t #0 [a::BIGINT]\n",
" Join INNER on=[(#1.0::BIGINT = #2.0::BIGINT)::BOOLEAN]\n",
" Get memory.main.v AS v #2 [c::BIGINT]\n",
" Get memory.main.u AS u #1 [b::BIGINT]\n",
)
);
}
#[test]
fn an_order_the_search_does_not_improve_on_is_left_exactly_as_it_was() {
let text = concat!(
"Join INNER on=[(#3.0::BIGINT = #1.0::BIGINT)::BOOLEAN]\n",
" Get memory.main.w AS w #3 [d::BIGINT]\n",
" Join INNER on=[(#1.0::BIGINT = #2.0::BIGINT)::BOOLEAN]\n",
" Get memory.main.u AS u #1 [b::BIGINT]\n",
" Get memory.main.v AS v #2 [c::BIGINT]\n",
);
assert_eq!(ordered(text), text);
}
#[test]
fn a_region_with_a_leaf_nobody_counted_is_left_alone() {
let text = concat!(
"Join INNER on=[(#1.0::BIGINT = #0.0::BIGINT)::BOOLEAN, ",
"(#1.0::BIGINT = #2.0::BIGINT)::BOOLEAN]\n",
" CrossProduct\n",
" Get memory.main.t AS t #0 [a::BIGINT]\n",
" Get memory.main.x AS x #2 [c::BIGINT]\n",
" Get memory.main.u AS u #1 [b::BIGINT]\n",
);
assert_eq!(ordered(text), text);
}
#[test]
fn a_condition_that_reads_one_leaf_stops_the_search() {
let text = concat!(
"Join INNER on=[(#1.0::BIGINT = #0.0::BIGINT)::BOOLEAN, ",
"(#2.0::BIGINT = #2.0::BIGINT)::BOOLEAN]\n",
" CrossProduct\n",
" Get memory.main.t AS t #0 [a::BIGINT]\n",
" Get memory.main.v AS v #2 [c::BIGINT]\n",
" Get memory.main.u AS u #1 [b::BIGINT]\n",
);
assert_eq!(ordered(text), text);
}
#[test]
fn an_outer_join_is_not_part_of_a_region() {
let text = concat!(
"Join INNER on=[(#1.0::BIGINT = #0.0::BIGINT)::BOOLEAN]\n",
" Join LEFT on=[(#0.0::BIGINT = #2.0::BIGINT)::BOOLEAN]\n",
" Get memory.main.t AS t #0 [a::BIGINT]\n",
" Get memory.main.v AS v #2 [c::BIGINT]\n",
" Get memory.main.u AS u #1 [b::BIGINT]\n",
);
assert_eq!(ordered(text), text);
}
#[test]
fn two_small_parts_with_no_condition_between_them_are_not_joined_to_each_other() {
assert_eq!(
ordered(concat!(
"Join INNER on=[(#0.0::BIGINT = #2.0::BIGINT)::BOOLEAN]\n",
" Join INNER on=[(#3.0::BIGINT = #0.0::BIGINT)::BOOLEAN, ",
"(#0.0::BIGINT = #1.0::BIGINT)::BOOLEAN]\n",
" CrossProduct\n",
" Get memory.main.w AS w #3 [d::BIGINT]\n",
" Get memory.main.u AS u #1 [b::BIGINT]\n",
" Get memory.main.t AS t #0 [a::BIGINT]\n",
" Get memory.main.v AS v #2 [c::BIGINT]\n",
)),
concat!(
"Join INNER on=[(#3.0::BIGINT = #0.0::BIGINT)::BOOLEAN]\n",
" Get memory.main.w AS w #3 [d::BIGINT]\n",
" Join INNER on=[(#0.0::BIGINT = #2.0::BIGINT)::BOOLEAN]\n",
" Get memory.main.v AS v #2 [c::BIGINT]\n",
" Join INNER on=[(#0.0::BIGINT = #1.0::BIGINT)::BOOLEAN]\n",
" Get memory.main.u AS u #1 [b::BIGINT]\n",
" Get memory.main.t AS t #0 [a::BIGINT]\n",
)
);
}
#[test]
fn a_region_with_no_path_between_two_parts_is_left_as_it_is() {
let text = concat!(
"CrossProduct\n",
" CrossProduct\n",
" Get memory.main.w AS w #3 [d::BIGINT]\n",
" Get memory.main.v AS v #2 [c::BIGINT]\n",
" Get memory.main.u AS u #1 [b::BIGINT]\n",
);
assert_eq!(ordered(text), text);
}
#[test]
fn an_order_that_is_cheaper_but_has_the_cross_products_it_already_had_is_not_taken() {
let text = concat!(
"CrossProduct\n",
" Join INNER on=[(#0.0::BIGINT = #1.0::BIGINT)::BOOLEAN]\n",
" Get memory.main.t AS t #0 [a::BIGINT]\n",
" Get memory.main.u AS u #1 [b::BIGINT]\n",
" Get memory.main.v AS v #2 [c::BIGINT]\n",
);
assert_eq!(ordered(text), text);
}
#[test]
fn a_region_under_something_else_is_reordered_and_what_is_above_it_is_rebuilt() {
assert_eq!(
ordered(concat!(
"Project #4 [#0.0::BIGINT AS a]\n",
" Join INNER on=[(#1.0::BIGINT = #0.0::BIGINT)::BOOLEAN, ",
"(#1.0::BIGINT = #2.0::BIGINT)::BOOLEAN]\n",
" CrossProduct\n",
" Get memory.main.t AS t #0 [a::BIGINT]\n",
" Get memory.main.v AS v #2 [c::BIGINT]\n",
" Get memory.main.u AS u #1 [b::BIGINT]\n",
)),
concat!(
"Project #4 [#0.0::BIGINT AS a]\n",
" Join INNER on=[(#1.0::BIGINT = #0.0::BIGINT)::BOOLEAN]\n",
" Get memory.main.t AS t #0 [a::BIGINT]\n",
" Join INNER on=[(#1.0::BIGINT = #2.0::BIGINT)::BOOLEAN]\n",
" Get memory.main.v AS v #2 [c::BIGINT]\n",
" Get memory.main.u AS u #1 [b::BIGINT]\n",
)
);
}
}