use alloc::string::String;
use alloc::vec::Vec;
use spg_sql::ast::{Expr, SelectItem, SelectStatement};
use crate::index_access::try_index_seek;
use spg_storage::{ColumnSchema, DataType, Row, Value};
use crate::{
CancelToken, Engine, EngineError, QueryResult, aggregate, expr_has_subquery, select_has_window,
};
pub(crate) fn build_index_suggestions(stmt: &SelectStatement, engine: &Engine) -> Vec<String> {
use alloc::collections::BTreeSet;
let mut seen: BTreeSet<(String, String)> = BTreeSet::new();
let mut out: Vec<String> = Vec::new();
let cat = engine.active_catalog();
let Some(from) = &stmt.from else {
return out;
};
let mut tables: Vec<String> = Vec::new();
tables.push(from.primary.name.clone());
for j in &from.joins {
tables.push(j.table.name.clone());
}
let mut col_refs: Vec<spg_sql::ast::ColumnName> = Vec::new();
if let Some(w) = &stmt.where_ {
collect_column_refs(w, &mut col_refs);
}
for j in &from.joins {
if let Some(on) = &j.on {
collect_column_refs(on, &mut col_refs);
}
}
for cn in &col_refs {
let owner: Option<String> = if let Some(q) = &cn.qualifier {
tables.iter().find(|t| t == &q).cloned()
} else {
tables.iter().find_map(|t| {
cat.get(t).and_then(|tbl| {
if tbl.schema().column_position(&cn.name).is_some() {
Some(t.clone())
} else {
None
}
})
})
};
let Some(owner) = owner else {
continue;
};
let Some(tbl) = cat.get(&owner) else {
continue;
};
let Some(col_pos) = tbl.schema().column_position(&cn.name) else {
continue;
};
let already_indexed = tbl.indices().iter().any(|i| {
matches!(i.kind, spg_storage::IndexKind::BTree(_))
&& i.column_position == col_pos
&& i.expression.is_none()
&& i.partial_predicate.is_none()
});
if already_indexed {
continue;
}
if seen.insert((owner.clone(), cn.name.clone())) {
out.push(alloc::format!(
"SUGGEST: CREATE INDEX ix_{}_{} ON {} ({})",
owner,
cn.name,
owner,
cn.name
));
}
}
let mut composite_eqs: alloc::collections::BTreeMap<
String,
alloc::collections::BTreeSet<String>,
> = alloc::collections::BTreeMap::new();
if let Some(w) = &stmt.where_ {
collect_and_eq_columns(w, &tables, cat, &mut composite_eqs);
}
for j in &from.joins {
if let Some(on) = &j.on {
collect_and_eq_columns(on, &tables, cat, &mut composite_eqs);
}
}
for (owner, cols) in composite_eqs {
if cols.len() < 2 {
continue;
}
let cols_vec: Vec<&String> = cols.iter().collect();
if let Some(tbl) = cat.get(&owner) {
let pos_to_name = |pos: usize| tbl.schema().columns.get(pos).map(|c| c.name.clone());
let already_in_index = tbl.indices().iter().any(|i| {
if !matches!(i.kind, spg_storage::IndexKind::BTree(_)) {
return false;
}
let mut all_cols: alloc::collections::BTreeSet<String> =
alloc::collections::BTreeSet::new();
if let Some(n) = pos_to_name(i.column_position) {
all_cols.insert(n);
}
for &extra in &i.extra_column_positions {
if let Some(c) = pos_to_name(extra) {
all_cols.insert(c);
}
}
cols.iter().all(|c| all_cols.contains(c))
});
let already_in_uc = tbl.schema().uniqueness_constraints.iter().any(|uc| {
let names: alloc::collections::BTreeSet<String> =
uc.columns.iter().filter_map(|&p| pos_to_name(p)).collect();
cols.iter().all(|c| names.contains(c))
});
if already_in_index || already_in_uc {
continue;
}
}
let cols_csv: Vec<String> = cols_vec.iter().map(|s| (*s).clone()).collect();
let suffix = cols_csv.join("_");
let body = cols_csv.join(", ");
out.push(alloc::format!(
"SUGGEST: CREATE INDEX ix_{owner}_{suffix} ON {owner} ({body})"
));
}
out
}
fn collect_and_eq_columns(
expr: &Expr,
tables: &[String],
cat: &spg_storage::Catalog,
out: &mut alloc::collections::BTreeMap<String, alloc::collections::BTreeSet<String>>,
) {
let mut stack: Vec<&Expr> = alloc::vec![expr];
while let Some(e) = stack.pop() {
match e {
Expr::Binary {
lhs,
op: spg_sql::ast::BinOp::And,
rhs,
} => {
stack.push(lhs);
stack.push(rhs);
}
Expr::Binary {
lhs,
op: spg_sql::ast::BinOp::Eq,
rhs,
} => {
let resolve = |e: &Expr| -> Option<(String, String)> {
if let Expr::Column(cn) = e {
let owner: Option<String> = if let Some(q) = &cn.qualifier {
tables.iter().find(|t| t == &q).cloned()
} else {
tables.iter().find_map(|t| {
cat.get(t).and_then(|tbl| {
if tbl.schema().column_position(&cn.name).is_some() {
Some(t.clone())
} else {
None
}
})
})
};
owner.map(|o| (o, cn.name.clone()))
} else {
None
}
};
let lhs_col = resolve(lhs);
let rhs_col = resolve(rhs);
if let (Some((t, c)), None) | (None, Some((t, c))) = (lhs_col, rhs_col) {
let mut entry = out.remove(&t).unwrap_or_default();
entry.insert(c);
out.insert(t, entry);
}
}
_ => {}
}
}
}
pub(crate) fn collect_column_refs(expr: &Expr, out: &mut Vec<spg_sql::ast::ColumnName>) {
match expr {
Expr::Column(cn) => out.push(cn.clone()),
Expr::FunctionCall { args, .. } => {
for a in args {
collect_column_refs(a, out);
}
}
Expr::Binary { lhs, rhs, .. } => {
collect_column_refs(lhs, out);
collect_column_refs(rhs, out);
}
Expr::Unary { expr: e, .. } => collect_column_refs(e, out),
_ => {}
}
}
struct PlanNode {
head: String,
attrs: Vec<String>,
children: Vec<PlanNode>,
no_arrow: bool,
cost: Option<(f64, f64, u64, u64)>,
actual: Option<(Option<f64>, u64)>,
}
impl PlanNode {
fn new(head: String) -> Self {
Self {
head,
attrs: Vec::new(),
children: Vec::new(),
no_arrow: false,
cost: None,
actual: None,
}
}
}
fn est_width(cols: &[ColumnSchema]) -> u64 {
cols.iter()
.map(|c| match c.ty {
DataType::SmallInt => 2,
DataType::Int | DataType::Date | DataType::Float => 4,
DataType::BigInt | DataType::Timestamp | DataType::Timestamptz | DataType::Money => 8,
DataType::Bool => 1,
DataType::Uuid => 16,
_ => 32,
})
.sum()
}
fn est_scan_rows(n: u64, where_: Option<&Expr>, eq_seek: bool) -> u64 {
match where_ {
None => n,
Some(_) if eq_seek => 1,
Some(w) => {
let frac = if matches!(
w,
Expr::Binary {
op: spg_sql::ast::BinOp::Eq,
..
}
) {
10
} else {
3
};
(n / frac).max(1)
}
}
}
fn render_pg_tree(node: &PlanNode, depth: usize, out: &mut Vec<String>) {
let head = if depth == 0 {
node.head.clone()
} else if node.no_arrow {
alloc::format!("{}{}", " ".repeat(6 * depth - 4), node.head)
} else {
alloc::format!("{}-> {}", " ".repeat(6 * depth - 6 + 2), node.head)
};
out.push(head);
let attr_pad = " ".repeat(6 * depth + 2);
for a in &node.attrs {
out.push(alloc::format!("{attr_pad}{a}"));
}
for c in &node.children {
render_pg_tree(c, depth + 1, out);
}
}
fn where_equi_join_conds<'w>(
from: &spg_sql::ast::FromClause,
where_: Option<&'w Expr>,
) -> alloc::vec::Vec<&'w Expr> {
let Some(w) = where_ else {
return alloc::vec::Vec::new();
};
if !from.joins.iter().all(|j| {
matches!(
j.kind,
spg_sql::ast::JoinKind::Inner | spg_sql::ast::JoinKind::Cross
)
}) {
return alloc::vec::Vec::new();
}
crate::reorder::split_and_conjunctions(w)
.into_iter()
.filter(|sub| equi_quals(sub).is_some())
.collect()
}
fn equi_quals(sub: &Expr) -> Option<(&str, &str)> {
let Expr::Binary {
lhs,
op: spg_sql::ast::BinOp::Eq,
rhs,
} = sub
else {
return None;
};
let (Expr::Column(a), Expr::Column(b)) = (lhs.as_ref(), rhs.as_ref()) else {
return None;
};
let (Some(qa), Some(qb)) = (a.qualifier.as_deref(), b.qualifier.as_deref()) else {
return None;
};
(!qa.eq_ignore_ascii_case(qb)).then_some((qa, qb))
}
fn promoted_key_for<'w>(
from: &spg_sql::ast::FromClause,
jidx: usize,
candidates: &[&'w Expr],
) -> Option<&'w Expr> {
let rel = |t: &spg_sql::ast::TableRef| t.alias.clone().unwrap_or_else(|| t.name.clone());
let peer = rel(&from.joins[jidx].table);
let mut left: alloc::vec::Vec<String> = alloc::vec![rel(&from.primary)];
left.extend(from.joins[..jidx].iter().map(|j| rel(&j.table)));
candidates.iter().copied().find(|sub| {
equi_quals(sub).is_some_and(|(qa, qb)| {
let names = |q: &str| left.iter().any(|l| l.eq_ignore_ascii_case(q));
(qa.eq_ignore_ascii_case(&peer) && names(qb))
|| (qb.eq_ignore_ascii_case(&peer) && names(qa))
})
})
}
fn without_conjuncts(where_: Option<&Expr>, drop: &[&Expr]) -> Option<Expr> {
if drop.is_empty() {
return None;
}
let w = where_?;
let dropped: alloc::collections::BTreeSet<usize> = drop
.iter()
.map(|e| core::ptr::from_ref::<Expr>(*e) as usize)
.collect();
crate::reorder::split_and_conjunctions(w)
.into_iter()
.filter(|c| !dropped.contains(&(core::ptr::from_ref::<Expr>(c) as usize)))
.cloned()
.reduce(|a, b| Expr::Binary {
lhs: alloc::boxed::Box::new(a),
op: spg_sql::ast::BinOp::And,
rhs: alloc::boxed::Box::new(b),
})
}
fn pg_cond(e: &Expr) -> String {
let s = alloc::format!("{e}");
if s.starts_with('(') && s.ends_with(')') {
s
} else {
alloc::format!("({s})")
}
}
fn split_index_cond<'a>(
engine: &Engine,
name: &str,
alias: &str,
where_: &'a Expr,
) -> (Option<&'a Expr>, Vec<&'a Expr>) {
fn flatten<'b>(e: &'b Expr, out: &mut Vec<&'b Expr>) {
if let Expr::Binary {
lhs,
op: spg_sql::ast::BinOp::And,
rhs,
} = e
{
flatten(lhs, out);
flatten(rhs, out);
} else {
out.push(e);
}
}
let mut conjuncts: Vec<&Expr> = Vec::new();
flatten(where_, &mut conjuncts);
let Some(table) = engine.active_catalog().get(name) else {
return (None, conjuncts);
};
let cols = &table.schema().columns;
let snap = engine.current_snapshot();
if conjuncts.len() == 1 {
return (Some(where_), Vec::new());
}
for (i, c) in conjuncts.iter().enumerate() {
if try_index_seek(c, cols, engine.active_catalog(), table, alias, &snap).is_some() {
let residual: Vec<&Expr> = conjuncts
.iter()
.enumerate()
.filter(|(j, _)| *j != i)
.map(|(_, e)| *e)
.collect();
return (Some(conjuncts[i]), residual);
}
}
(None, conjuncts)
}
fn index_name_for_cond(engine: &Engine, table: &str, alias: &str, cond: &Expr) -> Option<String> {
fn column_of<'a>(e: &'a Expr, out: &mut Vec<&'a spg_sql::ast::ColumnName>) {
match e {
Expr::Column(c) => out.push(c),
Expr::Binary { lhs, rhs, .. } => {
column_of(lhs, out);
column_of(rhs, out);
}
Expr::Unary { expr, .. } | Expr::Cast { expr, .. } => column_of(expr, out),
_ => {}
}
}
let mut refs: Vec<&spg_sql::ast::ColumnName> = Vec::new();
column_of(cond, &mut refs);
let t = engine.active_catalog().get(table)?;
let cols = &t.schema().columns;
for r in refs {
if let Some(q) = &r.qualifier
&& !q.eq_ignore_ascii_case(alias)
&& !q.eq_ignore_ascii_case(table)
{
continue;
}
let Some(pos) = cols
.iter()
.position(|c| c.name.eq_ignore_ascii_case(&r.name))
else {
continue;
};
if let Some(idx) = t.indices().iter().find(|i| {
i.column_position == pos && matches!(i.kind, spg_storage::IndexKind::BTree(_))
}) {
return Some(idx.name.clone());
}
}
None
}
fn pg_conjuncts(list: &[&Expr]) -> String {
match list.len() {
0 => String::new(),
1 => pg_cond(list[0]),
_ => {
let parts: Vec<String> = list.iter().map(|e| pg_cond(e)).collect();
alloc::format!("({})", parts.join(" AND "))
}
}
}
fn scan_node(
engine: &Engine,
name: &str,
alias: Option<&str>,
where_: Option<&Expr>,
cte_names: &[String],
index_only: bool,
) -> PlanNode {
let alias_sfx = match alias {
Some(a) if a != name => alloc::format!(" {a}"),
_ => String::new(),
};
if cte_names.iter().any(|c| c == name) {
let mut n = PlanNode::new(alloc::format!("CTE Scan on {name}{alias_sfx}"));
if let Some(w) = where_ {
n.attrs.push(alloc::format!("Filter: {}", pg_cond(w)));
}
return n;
}
if crate::partition::is_partition_parent(engine.active_catalog(), name) {
let mut app = PlanNode::new(String::from("Append"));
let kept = where_
.and_then(|_| engine.explain_partition_kept_children_by_where(name, where_))
.or_else(|| engine.explain_partition_kept_children_by_where(name, None));
if let Some(children) = kept {
for c in &children {
let mut sc = PlanNode::new(alloc::format!("Seq Scan on {c}"));
if let Some(w) = where_ {
sc.attrs.push(alloc::format!("Filter: {}", pg_cond(w)));
}
app.children.push(sc);
}
}
if app.children.is_empty() {
let mut sc = PlanNode::new(alloc::format!("Seq Scan on {name}{alias_sfx}"));
if let Some(w) = where_ {
sc.attrs.push(alloc::format!("Filter: {}", pg_cond(w)));
}
app.children.push(sc);
}
return app;
}
let seek = where_.and_then(|w| {
let table = engine.active_catalog().get(name)?;
let cols = &table.schema().columns;
let a = alias.unwrap_or(name);
try_index_seek(
w,
cols,
engine.active_catalog(),
table,
a,
&engine.current_snapshot(),
)
.map(|_| ())
});
let table_rows = engine
.active_catalog()
.get(name)
.map(|t| t.rows().len() as u64)
.unwrap_or(0);
let width = engine
.active_catalog()
.get(name)
.map(|t| est_width(&t.schema().columns))
.unwrap_or(8);
if seek.is_some() || index_only {
let a = alias.unwrap_or(name);
let split = where_.map(|w| split_index_cond(engine, name, a, w));
let idx_name = split
.as_ref()
.and_then(|(cond, _)| cond.as_ref().copied())
.or(where_)
.and_then(|c| index_name_for_cond(engine, name, a, c))
.or_else(|| {
engine.active_catalog().get(name).and_then(|t| {
t.indices()
.iter()
.find(|i| matches!(i.kind, spg_storage::IndexKind::BTree(_)))
.map(|i| i.name.clone())
})
})
.unwrap_or_else(|| alloc::format!("{name}_idx"));
let verb = if index_only {
"Index Only Scan using"
} else {
"Index Scan using"
};
let mut n = PlanNode::new(alloc::format!("{verb} {idx_name} on {name}{alias_sfx}"));
if let Some(w) = where_ {
let (cond, residual) = split.expect("computed alongside where_");
match cond {
Some(c) => {
n.attrs.push(alloc::format!("Index Cond: {}", pg_cond(c)));
if !residual.is_empty() {
n.attrs
.push(alloc::format!("Filter: {}", pg_conjuncts(&residual)));
}
}
None => n.attrs.push(alloc::format!("Index Cond: {}", pg_cond(w))),
}
}
let rows = est_scan_rows(table_rows, where_, true);
n.cost = Some((0.15, 0.15 + 8.0 + rows as f64 * 0.01, rows, width));
n
} else {
let mut n = PlanNode::new(alloc::format!("Seq Scan on {name}{alias_sfx}"));
let filtered = where_.is_some();
if let Some(w) = where_ {
n.attrs.push(alloc::format!("Filter: {}", pg_cond(w)));
}
let rows = est_scan_rows(table_rows, where_, false);
let total = 1.0
+ table_rows as f64 * 0.01
+ if filtered {
table_rows as f64 * 0.0025
} else {
0.0
};
n.cost = Some((0.0, total, rows, width));
n
}
}
fn child_cost(n: &PlanNode) -> (f64, f64, u64, u64) {
n.children
.first()
.and_then(|c| c.cost)
.unwrap_or((0.0, 0.0, 1, 8))
}
fn annotate_sort_method(node: &mut PlanNode, has_limit: bool) {
if node.head == "Sort"
&& !node.attrs.iter().any(|a| a.starts_with("Sort Method:"))
&& let Some(pos) = node.attrs.iter().position(|a| a.starts_with("Sort Key:"))
{
node.attrs.insert(
pos + 1,
alloc::string::String::from(if has_limit {
"Sort Method: top-N heapsort"
} else {
"Sort Method: quicksort"
}),
);
}
for c in &mut node.children {
annotate_sort_method(c, has_limit);
}
}
fn fill_actuals(
node: &mut PlanNode,
is_top: bool,
engine: &Engine,
result_rows: u64,
elapsed_ms: Option<f64>,
deltas: &alloc::collections::BTreeMap<String, u64>,
) {
fn scan_table(head: &str) -> Option<(&str, bool)> {
if let Some(r) = head.strip_prefix("Seq Scan on ") {
Some((r.split_whitespace().next().unwrap_or(r), true))
} else if let Some(r) = head.strip_prefix("CTE Scan on ") {
Some((r.split_whitespace().next().unwrap_or(r), false))
} else if let Some(r) = head
.strip_prefix("Index Scan using ")
.or_else(|| head.strip_prefix("Index Only Scan using "))
.and_then(|r| r.split_once(" on ").map(|(_, t)| t))
{
Some((r.split_whitespace().next().unwrap_or(r), false))
} else {
None
}
}
let head = node.head.clone();
let filtered = node.attrs.iter().any(|a| a.starts_with("Filter: "));
let live_rows = |t: &str| -> Option<u64> {
engine
.active_catalog()
.get(t)
.map(|tb| (tb.rows().len() as u64).saturating_sub(tb.dead_rows()))
};
if head.starts_with("Index Only Scan using ")
&& !node.attrs.iter().any(|a| a.starts_with("Heap Fetches:"))
{
node.attrs.push(String::from("Heap Fetches: 0"));
}
if is_top {
node.actual = Some((elapsed_ms, result_rows));
if node.children.is_empty()
&& filtered
&& let Some((t, is_seq)) = scan_table(&head)
&& is_seq
&& let Some(read) = live_rows(t)
&& read >= result_rows
{
node.attrs.push(alloc::format!(
"Rows Removed by Filter: {}",
read - result_rows
));
}
} else if let Some((table, is_seq)) = scan_table(&head) {
if !filtered {
if is_seq {
if let Some(n) = live_rows(table) {
node.actual = Some((None, n));
}
} else if let Some(&n) = deltas.get(table) {
node.actual = Some((None, n));
}
}
}
for c in &mut node.children {
fill_actuals(c, false, engine, result_rows, elapsed_ms, deltas);
}
}
fn scan_counter_snapshot(engine: &Engine) -> alloc::collections::BTreeMap<String, u64> {
use core::sync::atomic::Ordering;
let cat = engine.active_catalog();
let mut out = alloc::collections::BTreeMap::new();
for name in cat.table_names() {
if let Some(t) = cat.get(&name) {
let st = t.scan_stats();
let v =
st.seq_tup_read.load(Ordering::Relaxed) + st.idx_tup_fetch.load(Ordering::Relaxed);
out.insert(name, v);
}
}
out
}
enum Prop {
Str(String),
Bare(String),
List(Vec<String>),
}
fn node_props(node: &PlanNode, with_costs: bool, parent_rel: Option<&str>) -> Vec<(String, Prop)> {
let mut p: Vec<(String, Prop)> = Vec::new();
let mut push = |k: &str, v: Prop| p.push((String::from(k), v));
let head = node.head.as_str();
let (node_type, rel, idx) = if let Some(rest) = head.strip_prefix("Seq Scan on ") {
(
"Seq Scan",
Some(rest.split_whitespace().next().unwrap_or(rest)),
None,
)
} else if let Some(rest) = head.strip_prefix("Index Only Scan using ") {
let (i, r) = rest.split_once(" on ").unwrap_or((rest, ""));
(
"Index Only Scan",
Some(r.split_whitespace().next().unwrap_or(r)),
Some(i),
)
} else if let Some(rest) = head.strip_prefix("Index Scan using ") {
let (i, r) = rest.split_once(" on ").unwrap_or((rest, ""));
(
"Index Scan",
Some(r.split_whitespace().next().unwrap_or(r)),
Some(i),
)
} else if let Some(rest) = head.strip_prefix("CTE Scan on ") {
(
"CTE Scan",
Some(rest.split_whitespace().next().unwrap_or(rest)),
None,
)
} else if let Some(rest) = head.strip_prefix("Insert on ") {
("Insert", Some(rest), None)
} else if let Some(rest) = head.strip_prefix("Update on ") {
("Update", Some(rest), None)
} else if let Some(rest) = head.strip_prefix("Delete on ") {
("Delete", Some(rest), None)
} else if head.starts_with("CTE ") {
("CTE", None, None)
} else {
(head, None, None)
};
let is_agg = node_type == "HashAggregate" || node_type == "Aggregate";
push(
"Node Type",
Prop::Str(String::from(if is_agg { "Aggregate" } else { node_type })),
);
if let Some(pr) = parent_rel {
push("Parent Relationship", Prop::Str(String::from(pr)));
}
if is_agg {
let strategy = if node_type == "HashAggregate" {
"Hashed"
} else {
"Plain"
};
push("Strategy", Prop::Str(String::from(strategy)));
push("Partial Mode", Prop::Str(String::from("Simple")));
}
push("Parallel Aware", Prop::Bare(String::from("false")));
push("Async Capable", Prop::Bare(String::from("false")));
if let Some(i) = idx {
push("Scan Direction", Prop::Str(String::from("Forward")));
push("Index Name", Prop::Str(String::from(i)));
}
if let Some(r) = rel {
push("Relation Name", Prop::Str(String::from(r)));
push("Alias", Prop::Str(String::from(r)));
}
if with_costs && let Some((cs, ct, rows, width)) = node.cost {
push("Startup Cost", Prop::Bare(alloc::format!("{cs:.2}")));
push("Total Cost", Prop::Bare(alloc::format!("{ct:.2}")));
push("Plan Rows", Prop::Bare(alloc::format!("{rows}")));
push("Plan Width", Prop::Bare(alloc::format!("{width}")));
}
if let Some((time, rows)) = &node.actual {
if let Some(ms) = time {
push("Actual Startup Time", Prop::Bare(String::from("0.000")));
push("Actual Total Time", Prop::Bare(alloc::format!("{ms:.3}")));
}
push("Actual Rows", Prop::Bare(alloc::format!("{rows}.00")));
push("Actual Loops", Prop::Bare(String::from("1")));
}
push("Disabled", Prop::Bare(String::from("false")));
for a in &node.attrs {
let Some((k, v)) = a.split_once(": ") else {
continue;
};
if k == "Sort Key" || k == "Group Key" {
push(k, Prop::List(v.split(", ").map(String::from).collect()));
} else if k == "Rows Removed by Filter" {
push(k, Prop::Bare(String::from(v)));
} else {
push(k, Prop::Str(String::from(v)));
}
}
p
}
fn child_rel(i: usize) -> &'static str {
if i == 0 { "Outer" } else { "Inner" }
}
fn render_json_plan(node: &PlanNode, with_costs: bool) -> String {
fn obj(
node: &PlanNode,
with_costs: bool,
parent_rel: Option<&str>,
ind: usize,
out: &mut String,
) {
let pad = " ".repeat(ind);
let inner = " ".repeat(ind + 2);
out.push_str("{\n");
let props = node_props(node, with_costs, parent_rel);
let last = props.len().saturating_sub(1);
for (i, (k, v)) in props.iter().enumerate() {
out.push_str(&inner);
out.push_str(&json_string_lit(k));
out.push_str(": ");
match v {
Prop::Str(s) => out.push_str(&json_string_lit(s)),
Prop::Bare(s) => out.push_str(s),
Prop::List(items) => {
let its: Vec<String> = items.iter().map(|s| json_string_lit(s)).collect();
out.push_str(&alloc::format!("[{}]", its.join(", ")));
}
}
if i != last || !node.children.is_empty() {
out.push(',');
}
out.push('\n');
}
if !node.children.is_empty() {
out.push_str(&inner);
out.push_str("\"Plans\": [\n");
for (i, c) in node.children.iter().enumerate() {
out.push_str(&" ".repeat(ind + 4));
obj(c, with_costs, Some(child_rel(i)), ind + 4, out);
if i + 1 != node.children.len() {
out.push(',');
}
out.push('\n');
}
out.push_str(&inner);
out.push_str("]\n");
}
out.push_str(&pad);
out.push('}');
}
let mut out = String::from("[\n {\n \"Plan\": ");
obj(node, with_costs, None, 4, &mut out);
out.push_str("\n }\n]");
out
}
fn render_xml_plan(node: &PlanNode, with_costs: bool) -> String {
fn elem(
node: &PlanNode,
with_costs: bool,
parent_rel: Option<&str>,
ind: usize,
out: &mut String,
) {
let pad = " ".repeat(ind);
let inner = " ".repeat(ind + 2);
out.push_str(&alloc::format!("{pad}<Plan>\n"));
for (k, v) in node_props(node, with_costs, parent_rel) {
let tag = k.replace(' ', "-");
match v {
Prop::Str(s) => {
out.push_str(&alloc::format!(
"{inner}<{tag}>{}</{tag}>\n",
xml_escape(&s)
));
}
Prop::Bare(s) => out.push_str(&alloc::format!("{inner}<{tag}>{s}</{tag}>\n")),
Prop::List(items) => {
out.push_str(&alloc::format!("{inner}<{tag}>\n"));
for it in items {
out.push_str(&alloc::format!(
"{inner} <Item>{}</Item>\n",
xml_escape(&it)
));
}
out.push_str(&alloc::format!("{inner}</{tag}>\n"));
}
}
}
if !node.children.is_empty() {
out.push_str(&alloc::format!("{inner}<Plans>\n"));
for (i, c) in node.children.iter().enumerate() {
elem(c, with_costs, Some(child_rel(i)), ind + 4, out);
}
out.push_str(&alloc::format!("{inner}</Plans>\n"));
}
out.push_str(&alloc::format!("{pad}</Plan>\n"));
}
let mut out =
String::from("<explain xmlns=\"http://www.postgresql.org/2009/explain\">\n <Query>\n");
elem(node, with_costs, None, 4, &mut out);
out.push_str(" </Query>\n</explain>");
out
}
fn render_yaml_plan(node: &PlanNode, with_costs: bool) -> String {
fn map(
node: &PlanNode,
with_costs: bool,
parent_rel: Option<&str>,
ind: usize,
out: &mut String,
) {
let pad = " ".repeat(ind);
for (i, (k, v)) in node_props(node, with_costs, parent_rel).iter().enumerate() {
if i > 0 {
out.push_str(&pad);
}
match v {
Prop::Str(s) => out.push_str(&alloc::format!("{k}: {}\n", yaml_scalar(s))),
Prop::Bare(s) => out.push_str(&alloc::format!("{k}: {s}\n")),
Prop::List(items) => {
out.push_str(&alloc::format!("{k}: \n"));
for it in items {
out.push_str(&alloc::format!("{pad} - {}\n", yaml_scalar(it)));
}
}
}
}
if !node.children.is_empty() {
out.push_str(&alloc::format!("{pad}Plans: \n"));
for (i, c) in node.children.iter().enumerate() {
out.push_str(&alloc::format!("{pad} - "));
map(c, with_costs, Some(child_rel(i)), ind + 4, out);
}
}
}
let mut out = String::from("- Plan: \n ");
map(node, with_costs, None, 4, &mut out);
out
}
fn build_plan_tree(stmt: &SelectStatement, engine: &Engine) -> PlanNode {
let cte_names: Vec<String> = stmt.ctes.iter().map(|c| c.name.clone()).collect();
if !stmt.unions.is_empty() {
let mut root = PlanNode::new(String::from("Append"));
let mut first = stmt.clone();
first.unions = Vec::new();
root.children.push(build_plan_tree(&first, engine));
for (_kind, peer) in &stmt.unions {
root.children.push(build_plan_tree(peer, engine));
}
let mut total = 0.0f64;
let mut rows = 0u64;
let mut width = 8u64;
for c in &root.children {
if let Some((_, ct, cr, cw)) = c.cost {
total += ct;
rows += cr;
width = width.max(cw);
}
}
root.cost = Some((0.0, total + rows as f64 * 0.0025, rows, width));
return root;
}
let mut node = match &stmt.from {
None => {
let mut r = PlanNode::new(String::from("Result"));
r.cost = Some((0.0, 0.01, 1, 4));
r
}
Some(from) => {
let where_eq = where_equi_join_conds(from, stmt.where_.as_ref());
let promoted_all: alloc::vec::Vec<&Expr> = (0..from.joins.len())
.filter(|&i| {
!from.joins[i].on.as_ref().is_some_and(|on| {
matches!(on, Expr::Binary { op, .. } if matches!(op, spg_sql::ast::BinOp::Eq))
})
})
.filter_map(|i| promoted_key_for(from, i, &where_eq))
.collect();
let scan_where = without_conjuncts(stmt.where_.as_ref(), &promoted_all);
let mut left = scan_node(
engine,
&from.primary.name,
from.primary.alias.as_deref(),
scan_where.as_ref().or(if promoted_all.is_empty() {
stmt.where_.as_ref()
} else {
None
}),
&cte_names,
engine.stmt_takes_index_only_scan(stmt),
);
for (jidx, j) in from.joins.iter().enumerate() {
let right = scan_node(
engine,
&j.table.name,
j.table.alias.as_deref(),
None,
&cte_names,
false,
);
let (verb, hashable) = match j.kind {
spg_sql::ast::JoinKind::Inner => ("", true),
spg_sql::ast::JoinKind::Left => (" Left", true),
spg_sql::ast::JoinKind::Right => (" Right", true),
spg_sql::ast::JoinKind::FullOuter => (" Full", true),
spg_sql::ast::JoinKind::Semi => (" Semi", true),
spg_sql::ast::JoinKind::Cross => ("", true),
};
let is_eq_join = j
.on
.as_ref()
.is_some_and(|on| matches!(on, Expr::Binary { op, .. } if matches!(op, spg_sql::ast::BinOp::Eq)));
let promoted = if is_eq_join {
None
} else {
promoted_key_for(from, jidx, &where_eq)
};
let mut jn = if (is_eq_join || promoted.is_some()) && hashable {
let mut jn = PlanNode::new(alloc::format!("Hash Join{verb}"));
jn.head = alloc::format!("Hash{verb} Join");
if let Some(on) = j.on.as_ref().or(promoted) {
jn.attrs.push(alloc::format!("Hash Cond: {}", pg_cond(on)));
}
let mut hash = PlanNode::new(String::from("Hash"));
let (_, rt, rr, rw) = right.cost.unwrap_or((0.0, 0.0, 1, 8));
hash.cost = Some((rt, rt + rr as f64 * 0.01, rr, rw));
hash.children.push(right);
let (_, lt, lr, lw) = left.cost.unwrap_or((0.0, 0.0, 1, 8));
let (hs, ht, hr, hw) = hash.cost.unwrap_or((0.0, 0.0, 1, 8));
let _ = hs;
jn.cost = Some((ht, ht + lt + (lr + hr) as f64 * 0.01, lr.max(hr), lw + hw));
jn.children.push(left);
jn.children.push(hash);
jn
} else {
let mut jn = PlanNode::new(alloc::format!("Nested Loop{verb}"));
if let Some(on) = &j.on {
jn.attrs
.push(alloc::format!("Join Filter: {}", pg_cond(on)));
}
let (_, lt, lr, lw) = left.cost.unwrap_or((0.0, 0.0, 1, 8));
let (_, rt, rr, rw) = right.cost.unwrap_or((0.0, 0.0, 1, 8));
jn.cost = Some((0.0, lt + lr as f64 * rt.max(0.01), lr * rr.max(1), lw + rw));
jn.children.push(left);
jn.children.push(right);
jn
};
if expr_has_subquery(
stmt.where_
.as_ref()
.unwrap_or(&Expr::Literal(spg_sql::ast::Literal::Null)),
) {
let _ = &mut jn;
}
left = jn;
}
left
}
};
if select_has_window(stmt) {
let mut w = PlanNode::new(String::from("WindowAgg"));
w.children.push(node);
let (cs, ct, cr, cw) = child_cost(&w);
w.cost = Some((cs, ct + cr as f64 * 0.01, cr, cw + 8));
node = w;
}
if aggregate::uses_aggregate(stmt) || stmt.group_by.is_some() {
let mut agg = if let Some(gs) = &stmt.group_by {
let mut a = PlanNode::new(String::from("HashAggregate"));
let keys: Vec<String> = gs.iter().map(|g| alloc::format!("{g}")).collect();
a.attrs
.push(alloc::format!("Group Key: {}", keys.join(", ")));
a
} else {
PlanNode::new(String::from("Aggregate"))
};
if let Some(h) = &stmt.having {
agg.attrs.push(alloc::format!("Filter: {}", pg_cond(h)));
}
agg.children.push(node);
let (_, ct, cr, cw) = child_cost(&agg);
let out_rows = if stmt.group_by.is_some() {
(cr / 10).max(1)
} else {
1
};
agg.cost = Some((ct, ct + cr as f64 * 0.0025, out_rows, cw.min(16)));
node = agg;
} else if stmt.distinct {
let mut d = PlanNode::new(String::from("HashAggregate"));
let keys: Vec<String> = stmt
.items
.iter()
.map(|it| match it {
SelectItem::Wildcard => String::from("*"),
SelectItem::Expr { expr, .. } => alloc::format!("{expr}"),
other => alloc::format!("{other:?}"),
})
.collect();
d.attrs
.push(alloc::format!("Group Key: {}", keys.join(", ")));
d.children.push(node);
let (_, ct, cr, cw) = child_cost(&d);
d.cost = Some((ct, ct + cr as f64 * 0.0025, (cr / 10).max(1), cw));
node = d;
}
if !stmt.order_by.is_empty() {
let mut s = PlanNode::new(String::from("Sort"));
let keys: Vec<String> = stmt
.order_by
.iter()
.map(|o| {
if o.desc {
alloc::format!("{} DESC", o.expr)
} else {
alloc::format!("{}", o.expr)
}
})
.collect();
s.attrs
.push(alloc::format!("Sort Key: {}", keys.join(", ")));
s.children.push(node);
let (_, ct, cr, cw) = child_cost(&s);
let sort_cost = ct + cr as f64 * 0.02;
s.cost = Some((sort_cost, sort_cost + cr as f64 * 0.01, cr, cw));
node = s;
}
if stmt.limit.is_some() || stmt.offset.is_some() {
let mut l = PlanNode::new(String::from("Limit"));
l.children.push(node);
let (cs, ct, cr, cw) = child_cost(&l);
let lim = match &stmt.limit {
Some(spg_sql::ast::LimitExpr::Literal(n)) => u64::from(*n).min(cr),
_ => cr,
};
l.cost = Some((cs, ct, lim, cw));
node = l;
}
for cte in &stmt.ctes {
let label = if cte.recursive {
alloc::format!("CTE {} (recursive)", cte.name)
} else {
alloc::format!("CTE {}", cte.name)
};
let mut block = PlanNode::new(label);
block.no_arrow = true;
match &cte.body {
spg_sql::ast::CteBody::Select(s) => {
block.children.push(build_plan_tree(s, engine));
}
spg_sql::ast::CteBody::Insert(s) => {
block
.children
.push(PlanNode::new(alloc::format!("Insert on {}", s.table)));
}
spg_sql::ast::CteBody::Update(s) => {
block
.children
.push(PlanNode::new(alloc::format!("Update on {}", s.table)));
}
spg_sql::ast::CteBody::Delete(s) => {
block
.children
.push(PlanNode::new(alloc::format!("Delete on {}", s.table)));
}
spg_sql::ast::CteBody::Merge(s) => {
block
.children
.push(PlanNode::new(alloc::format!("Merge on {}", s.target)));
}
}
node.children.insert(0, block);
}
node
}
pub(crate) fn explain_select(
stmt: &SelectStatement,
engine: &Engine,
depth: usize,
out: &mut Vec<String>,
) {
let tree = build_plan_tree(stmt, engine);
render_pg_tree(&tree, depth, out);
}
pub(crate) fn explain_select_costed(
stmt: &SelectStatement,
engine: &Engine,
with_costs: bool,
out: &mut Vec<String>,
) {
let tree = build_plan_tree(stmt, engine);
render_costed(&tree, with_costs, out);
}
fn render_costed(tree: &PlanNode, with_costs: bool, out: &mut Vec<String>) {
fn walk(node: &PlanNode, depth: usize, out: &mut Vec<String>, with_costs: bool) {
let mut head = if depth == 0 {
node.head.clone()
} else if node.no_arrow {
alloc::format!("{}{}", " ".repeat(6 * depth - 4), node.head)
} else {
alloc::format!("{}-> {}", " ".repeat(6 * depth - 6 + 2), node.head)
};
if with_costs && let Some((cs, ct, rows, width)) = node.cost {
head.push_str(&alloc::format!(
" (cost={cs:.2}..{ct:.2} rows={rows} width={width})"
));
}
if let Some((t, rows)) = node.actual {
match t {
Some(ms) => head.push_str(&alloc::format!(
" (actual time=0.000..{ms:.3} rows={rows}.00 loops=1)"
)),
None => head.push_str(&alloc::format!(" (actual rows={rows}.00 loops=1)")),
}
}
out.push(head);
let attr_pad = " ".repeat(6 * depth + 2);
for a in &node.attrs {
out.push(alloc::format!("{attr_pad}{a}"));
}
for c in &node.children {
walk(c, depth + 1, out, with_costs);
}
}
walk(tree, 0, out, with_costs);
}
impl Engine {
fn dml_plan_tree(&self, inner: &spg_sql::ast::Statement) -> Option<PlanNode> {
match inner {
spg_sql::ast::Statement::Insert(i) => {
let mut root = PlanNode::new(alloc::format!("Insert on {}", i.table));
let child = match &i.select_source {
Some(src) => build_plan_tree(src, self),
None => {
let mut r = PlanNode::new(String::from("Result"));
r.cost = Some((0.0, 0.01, i.rows.len().max(1) as u64, 8));
r
}
};
root.cost = child.cost.map(|(_, ct, _, _)| (0.0, ct, 0, 0));
root.children.push(child);
Some(root)
}
spg_sql::ast::Statement::Update(u) => {
let mut root = PlanNode::new(alloc::format!("Update on {}", u.table));
let child = scan_node(self, &u.table, None, u.where_.as_ref(), &[], false);
root.cost = child.cost.map(|(cs, ct, _, _)| (cs, ct, 0, 0));
root.children.push(child);
Some(root)
}
spg_sql::ast::Statement::Delete(d) => {
let mut root = PlanNode::new(alloc::format!("Delete on {}", d.table));
let child = scan_node(self, &d.table, None, d.where_.as_ref(), &[], false);
root.cost = child.cost.map(|(cs, ct, _, _)| (cs, ct, 0, 0));
root.children.push(child);
Some(root)
}
_ => None,
}
}
}
impl Engine {
pub(crate) fn exec_explain_analyze_dml(
&mut self,
e: &spg_sql::ast::ExplainStatement,
cancel: CancelToken<'_>,
) -> Result<QueryResult, EngineError> {
let Some(mut tree) = self.dml_plan_tree(&e.inner) else {
return Err(EngineError::Unsupported(String::from(
"EXPLAIN ANALYZE body must be INSERT / UPDATE / DELETE",
)));
};
let started = self.clock.map(|f| f());
let res = self.dispatch_stmt_inner((*e.inner).clone(), cancel)?;
let elapsed_micros = match (self.clock, started) {
(Some(f), Some(s)) => Some(f().saturating_sub(s)),
_ => None,
};
let affected = match &res {
QueryResult::CommandOk { affected, .. } => *affected as u64,
QueryResult::Rows { rows, .. } => rows.len() as u64,
};
let show_time = !e.timing_off && !self.env_cfg().explain_no_costs;
let elapsed_ms = if show_time {
elapsed_micros.map(|us| us as f64 / 1000.0)
} else {
None
};
tree.actual = Some((elapsed_ms, 0));
for child in &mut tree.children {
child.actual = Some((None, affected));
}
let mut lines = Vec::<String>::new();
render_costed(&tree, !e.costs_off, &mut lines);
if !e.summary_off
&& !self.env_cfg().explain_no_costs
&& let Some(us) = elapsed_micros
{
let ms = us as f64 / 1000.0;
lines.push(alloc::format!("Execution Time: {ms:.3} ms"));
}
let columns = alloc::vec![ColumnSchema::new("QUERY PLAN", DataType::Text, false)];
let rows: Vec<Row<'static>> = lines
.into_iter()
.map(|l| Row::new(alloc::vec![Value::text(l)]))
.collect();
Ok(QueryResult::Rows { columns, rows })
}
#[allow(clippy::format_push_string)]
pub(crate) fn exec_explain(
&self,
e: &spg_sql::ast::ExplainStatement,
cancel: CancelToken<'_>,
) -> Result<QueryResult, EngineError> {
let mut lines = Vec::<String>::new();
let mut plan_tree: Option<PlanNode>;
let sel: Option<&SelectStatement> = match &*e.inner {
spg_sql::ast::Statement::Select(s) => {
let tree = build_plan_tree(s, self);
render_costed(&tree, !e.costs_off, &mut lines);
plan_tree = Some(tree);
Some(s)
}
dml @ (spg_sql::ast::Statement::Insert(_)
| spg_sql::ast::Statement::Update(_)
| spg_sql::ast::Statement::Delete(_)) => {
let root = self.dml_plan_tree(dml).expect("DML arm builds a tree");
render_costed(&root, !e.costs_off, &mut lines);
plan_tree = Some(root);
None
}
other => {
plan_tree = None;
let _ = &plan_tree;
return Err(EngineError::Unsupported(alloc::format!(
"EXPLAIN body must be SELECT / INSERT / UPDATE / DELETE, got {other:?}"
)));
}
};
if e.suggest {
if let Some(sel) = sel {
let suggestions = build_index_suggestions(sel, self);
for s in suggestions {
lines.push(s);
}
}
} else if e.analyze {
let Some(sel) = sel else {
return Err(EngineError::Unsupported(String::from(
"EXPLAIN ANALYZE on INSERT/UPDATE/DELETE cannot run on the read-only path",
)));
};
let before = scan_counter_snapshot(self);
let started = self.clock.map(|f| f());
let exec_rows =
self.execute_readonly_select_streaming_prepared(sel, cancel, |item| {
let _ = matches!(item, crate::StreamItem::Row(_));
Ok(())
})?;
let elapsed_micros = match (self.clock, started) {
(Some(f), Some(s)) => Some(f().saturating_sub(s)),
_ => None,
};
let after = scan_counter_snapshot(self);
let mut deltas: alloc::collections::BTreeMap<String, u64> =
alloc::collections::BTreeMap::new();
for (k, v) in &after {
let d = v.saturating_sub(before.get(k).copied().unwrap_or(0));
if d > 0 {
deltas.insert(k.clone(), d);
}
}
let row_count = exec_rows;
let show_time = !e.timing_off && !self.env_cfg().explain_no_costs;
let elapsed_ms = if show_time {
elapsed_micros.map(|us| us as f64 / 1000.0)
} else {
None
};
if let Some(tree) = &mut plan_tree {
fill_actuals(tree, true, self, row_count as u64, elapsed_ms, &deltas);
annotate_sort_method(tree, sel.limit.is_some());
lines.clear();
render_costed(tree, !e.costs_off, &mut lines);
}
if !e.summary_off
&& !self.env_cfg().explain_no_costs
&& let Some(us) = elapsed_micros
{
let ms = us as f64 / 1000.0;
lines.push(alloc::format!("Execution Time: {ms:.3} ms"));
}
if e.buffers {
let cold_rows: u64 = 0;
let hot_rows: u64 = row_count as u64;
let total_rows = hot_rows.saturating_add(cold_rows);
let ratio = if total_rows == 0 {
alloc::string::String::from("n/a")
} else {
let ratio_x10000 = (hot_rows.saturating_mul(10_000)) / total_rows;
alloc::format!("{}.{:02}", ratio_x10000 / 100, ratio_x10000 % 100)
};
lines.push(alloc::format!(
"Buffers: hot_rows={hot_rows} cold_rows={cold_rows} cache_hit_ratio={ratio}"
));
}
}
if e.settings {
let mut diverged: Vec<alloc::string::String> = Vec::new();
for key in [
"default_text_search_config",
"statement_timeout",
"default_transaction_isolation",
"search_path",
] {
if let Some(v) = self.session_param(key) {
diverged.push(alloc::format!("{key}={v}"));
}
}
if diverged.is_empty() {
lines.push("Settings: (no overrides)".into());
} else {
lines.push(alloc::format!("Settings: {}", diverged.join(", ")));
}
}
if e.wal {
lines.push("WAL: records=0 bytes=0 fpi=0".into());
}
let columns = alloc::vec![ColumnSchema::new("QUERY PLAN", DataType::Text, false)];
let rows: Vec<Row<'static>> = match e.format {
spg_sql::ast::ExplainFormat::Text => lines
.into_iter()
.map(|l| Row::new(alloc::vec![Value::text(l)]))
.collect(),
spg_sql::ast::ExplainFormat::Json => {
let body = match &plan_tree {
Some(tree) => render_json_plan(tree, !e.costs_off),
None => {
let mut b = alloc::string::String::from("[");
for (i, l) in lines.iter().enumerate() {
if i > 0 {
b.push_str(", ");
}
b.push_str("{\"Plan Line\": ");
b.push_str(&json_string_lit(l));
b.push('}');
}
b.push(']');
b
}
};
alloc::vec![Row::new(alloc::vec![Value::text(body)])]
}
spg_sql::ast::ExplainFormat::Xml => {
let body = match &plan_tree {
Some(tree) => render_xml_plan(tree, !e.costs_off),
None => {
let mut b = alloc::string::String::from(
"<explain xmlns=\"http://www.postgresql.org/2009/explain\">",
);
for l in &lines {
b.push_str("<line>");
b.push_str(&xml_escape(l));
b.push_str("</line>");
}
b.push_str("</explain>");
b
}
};
alloc::vec![Row::new(alloc::vec![Value::text(body)])]
}
spg_sql::ast::ExplainFormat::Yaml => {
let body = match &plan_tree {
Some(tree) => render_yaml_plan(tree, !e.costs_off),
None => {
let mut b = alloc::string::String::from("- Plan:\n");
for l in &lines {
b.push_str(" - ");
b.push_str(&yaml_scalar(l));
b.push('\n');
}
b
}
};
alloc::vec![Row::new(alloc::vec![Value::text(body)])]
}
};
Ok(QueryResult::Rows { columns, rows })
}
}
fn json_string_lit(s: &str) -> alloc::string::String {
let mut out = alloc::string::String::with_capacity(s.len() + 2);
out.push('"');
for c in s.chars() {
match c {
'"' => out.push_str("\\\""),
'\\' => out.push_str("\\\\"),
'\n' => out.push_str("\\n"),
'\r' => out.push_str("\\r"),
'\t' => out.push_str("\\t"),
c if (c as u32) < 0x20 => {
out.push_str(&alloc::format!("\\u{:04x}", c as u32));
}
c => out.push(c),
}
}
out.push('"');
out
}
fn xml_escape(s: &str) -> alloc::string::String {
let mut out = alloc::string::String::with_capacity(s.len());
for c in s.chars() {
match c {
'&' => out.push_str("&"),
'<' => out.push_str("<"),
'>' => out.push_str(">"),
'"' => out.push_str("""),
'\'' => out.push_str("'"),
c => out.push(c),
}
}
out
}
fn yaml_scalar(s: &str) -> alloc::string::String {
json_string_lit(s)
}