use alloc::collections::BTreeMap;
use alloc::format;
use alloc::string::{String, ToString};
use alloc::vec::Vec;
use spg_sql::ast::{
BinOp, ColumnName, Expr, FromClause, FromJoin, JoinKind, Literal, SelectItem, SelectStatement,
};
use spg_storage::Value;
use crate::cancel::CancelToken;
use crate::reorder::split_and_conjunctions;
use crate::{Engine, EngineError, QueryResult};
impl Engine {
pub(crate) fn try_fold_inner_joins(
&self,
stmt: &SelectStatement,
cancel: CancelToken<'_>,
) -> Result<Option<SelectStatement>, EngineError> {
let Some(from) = &stmt.from else {
return Ok(None);
};
if from.joins.is_empty() {
return Ok(None);
}
for j in &from.joins {
if !matches!(j.kind, JoinKind::Inner) {
return Ok(None);
}
}
if !stmt.order_by.is_empty()
|| stmt.limit.is_some()
|| stmt.offset.is_some()
|| stmt.having.is_some()
|| stmt.distinct
|| !stmt.unions.is_empty()
|| stmt.limit_with_ties
{
return Ok(None);
}
if statement_has_unresolved_subquery(stmt) {
return Ok(None);
}
for j in &from.joins {
if j.table.lateral_subquery.is_some()
|| j.table.unnest_expr.is_some()
|| j.table.generate_series_args.is_some()
|| j.table.as_of_segment.is_some()
{
return Ok(None);
}
}
if from.primary.lateral_subquery.is_some()
|| from.primary.unnest_expr.is_some()
|| from.primary.generate_series_args.is_some()
|| from.primary.as_of_segment.is_some()
{
return Ok(None);
}
let primary_alias = table_alias(&from.primary).to_string();
let mut alias_to_idx: BTreeMap<String, usize> = BTreeMap::new();
alias_to_idx.insert(primary_alias.clone(), 0);
for (i, j) in from.joins.iter().enumerate() {
let a = table_alias(&j.table).to_string();
if alias_to_idx.insert(a, i + 1).is_some() {
return Ok(None);
}
}
let mut planned: Vec<FoldPlan> = Vec::new();
let mut folded_join_indices: Vec<usize> = Vec::new();
for (i, j) in from.joins.iter().enumerate() {
let Some(on) = j.on.as_ref() else {
continue;
};
let inner_alias = table_alias(&j.table);
let Some((outer_expr, inner_col_name)) = analyse_on_eq(on, inner_alias, &alias_to_idx)
else {
continue;
};
if !self.column_is_single_unique_public(&j.table.name, &inner_col_name) {
continue;
}
if alias_referenced_elsewhere(stmt, &from.joins, i, inner_alias) {
continue;
}
let inner_preds: Vec<Expr> = if let Some(w) = &stmt.where_ {
split_and_conjunctions(w)
.into_iter()
.filter(|p| {
let refs = referenced_aliases(p, &alias_to_idx);
refs.len() == 1 && refs.iter().next().copied() == Some(i + 1)
})
.cloned()
.collect()
} else {
Vec::new()
};
if inner_preds.is_empty() {
continue;
}
let pk_values = match self.collect_unique_keys(
&j.table.name,
inner_alias,
&inner_col_name,
&inner_preds,
cancel,
) {
Ok(v) => v,
Err(_) => continue,
};
planned.push(FoldPlan {
join_idx: i,
outer_expr,
pk_values,
});
folded_join_indices.push(i);
}
if planned.is_empty() {
return Ok(None);
}
if folded_join_indices.len() != from.joins.len() {
return Ok(None);
}
let mut new_stmt = stmt.clone();
let new_from = build_folded_from(from, &folded_join_indices);
new_stmt.from = Some(new_from);
new_stmt.where_ = rewrite_where(stmt.where_.as_ref(), &alias_to_idx, &planned);
Ok(Some(new_stmt))
}
}
struct FoldPlan {
join_idx: usize,
outer_expr: Expr,
pk_values: Vec<Value<'static>>,
}
impl Engine {
pub(crate) fn try_eliminate_redundant_left_joins(
&self,
stmt: &SelectStatement,
) -> Option<SelectStatement> {
let from = stmt.from.as_ref()?;
if from.joins.is_empty() {
return None;
}
if !from.joins.iter().any(|j| matches!(j.kind, JoinKind::Left)) {
return None;
}
if from.primary.lateral_subquery.is_some()
|| from.primary.unnest_expr.is_some()
|| from.primary.generate_series_args.is_some()
|| from.primary.as_of_segment.is_some()
{
return None;
}
for j in &from.joins {
if j.table.lateral_subquery.is_some()
|| j.table.unnest_expr.is_some()
|| j.table.generate_series_args.is_some()
|| j.table.as_of_segment.is_some()
{
return None;
}
}
let primary_alias = table_alias(&from.primary).to_string();
let mut alias_to_idx: BTreeMap<String, usize> = BTreeMap::new();
alias_to_idx.insert(primary_alias.clone(), 0);
for (i, j) in from.joins.iter().enumerate() {
let a = table_alias(&j.table).to_string();
if alias_to_idx.insert(a, i + 1).is_some() {
return None;
}
}
let mut drop_indices: Vec<usize> = Vec::new();
for (i, j) in from.joins.iter().enumerate() {
if !matches!(j.kind, JoinKind::Left) {
continue;
}
let on = j.on.as_ref()?;
let inner_alias = table_alias(&j.table);
let (_outer_expr, inner_col_name) = analyse_on_eq(on, inner_alias, &alias_to_idx)?;
if !self.column_is_single_unique_public(&j.table.name, &inner_col_name) {
continue;
}
if alias_referenced_strictly_elsewhere(stmt, &from.joins, i, inner_alias) {
continue;
}
drop_indices.push(i);
}
if drop_indices.is_empty() {
return None;
}
let new_from = build_folded_from(from, &drop_indices);
let mut new_stmt = stmt.clone();
new_stmt.from = Some(new_from);
Some(new_stmt)
}
}
fn statement_has_unresolved_subquery(stmt: &SelectStatement) -> bool {
let mut hit = false;
let mut visit = |e: &Expr| {
walk_expr(e, &mut |item| match item {
WalkItem::Subquery(_) | WalkItem::Unknown => hit = true,
WalkItem::Column(_) => {}
});
};
if let Some(w) = &stmt.where_ {
visit(w);
}
for it in &stmt.items {
if let SelectItem::Expr { expr, .. } = it {
visit(expr);
}
}
for o in &stmt.order_by {
visit(&o.expr);
}
if let Some(g) = &stmt.group_by {
for e in g {
visit(e);
}
}
if let Some(h) = &stmt.having {
visit(h);
}
if let Some(from) = &stmt.from {
for j in &from.joins {
if let Some(on) = &j.on {
visit(on);
}
}
}
hit
}
fn table_alias(t: &spg_sql::ast::TableRef) -> &str {
t.alias.as_deref().unwrap_or(&t.name)
}
fn analyse_on_eq(
on: &Expr,
inner_alias: &str,
alias_to_idx: &BTreeMap<String, usize>,
) -> Option<(Expr, String)> {
let Expr::Binary {
op: BinOp::Eq,
lhs,
rhs,
} = on
else {
return None;
};
let (outer_side, inner_side) = if is_qualified_with(lhs, inner_alias) {
(rhs.as_ref(), lhs.as_ref())
} else if is_qualified_with(rhs, inner_alias) {
(lhs.as_ref(), rhs.as_ref())
} else {
return None;
};
let Expr::Column(ColumnName {
qualifier: _,
name: inner_col,
}) = inner_side
else {
return None;
};
let refs = referenced_aliases(outer_side, alias_to_idx);
let inner_idx = *alias_to_idx.get(inner_alias)?;
if refs.contains(&inner_idx) {
return None;
}
if refs.is_empty() {
return None;
}
Some((outer_side.clone(), inner_col.clone()))
}
fn is_qualified_with(e: &Expr, alias: &str) -> bool {
matches!(
e,
Expr::Column(ColumnName {
qualifier: Some(q),
..
}) if q == alias
)
}
fn referenced_aliases(
e: &Expr,
alias_to_idx: &BTreeMap<String, usize>,
) -> alloc::collections::BTreeSet<usize> {
let mut state = WalkState {
out: alloc::collections::BTreeSet::new(),
unknown: false,
};
walk_expr(e, &mut |kind| match kind {
WalkItem::Column(c) => match &c.qualifier {
Some(q) => match alias_to_idx.get(q) {
Some(&i) => {
state.out.insert(i);
}
None => state.unknown = true,
},
None => state.unknown = true,
},
WalkItem::Subquery(_) | WalkItem::Unknown => state.unknown = true,
});
if state.unknown {
state.out.clear();
for i in 0..alias_to_idx.len() {
state.out.insert(i);
}
}
state.out
}
struct WalkState {
out: alloc::collections::BTreeSet<usize>,
unknown: bool,
}
enum WalkItem<'a> {
Column(&'a ColumnName),
Subquery(&'a SelectStatement),
Unknown,
}
fn walk_expr<F>(e: &Expr, f: &mut F)
where
F: FnMut(WalkItem<'_>),
{
match e {
Expr::Column(c) => f(WalkItem::Column(c)),
Expr::Literal(_) | Expr::Placeholder(_) => {}
Expr::Binary { lhs, rhs, .. } => {
walk_expr(lhs, f);
walk_expr(rhs, f);
}
Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
walk_expr(expr, f)
}
Expr::FunctionCall { args, .. } => {
for a in args {
walk_expr(a, f);
}
}
Expr::Like { expr, pattern, .. } => {
walk_expr(expr, f);
walk_expr(pattern, f);
}
Expr::InList { expr: e2, list, .. } => {
walk_expr(e2, f);
for li in list {
walk_expr(li, f);
}
}
Expr::ScalarSubquery(s) => f(WalkItem::Subquery(s)),
Expr::Exists { subquery, .. } => f(WalkItem::Subquery(subquery)),
Expr::InSubquery {
expr: e2, subquery, ..
} => {
walk_expr(e2, f);
f(WalkItem::Subquery(subquery));
}
_ => f(WalkItem::Unknown),
}
}
fn alias_referenced_elsewhere(
stmt: &SelectStatement,
joins: &[FromJoin],
skip_idx: usize,
alias: &str,
) -> bool {
for it in &stmt.items {
match it {
SelectItem::Wildcard => continue,
SelectItem::Expr { expr, .. } => {
if expr_references_alias(expr, alias) {
return true;
}
}
}
}
for o in &stmt.order_by {
if expr_references_alias(&o.expr, alias) {
return true;
}
}
if let Some(g) = &stmt.group_by {
for e in g {
if expr_references_alias(e, alias) {
return true;
}
}
}
if let Some(h) = &stmt.having
&& expr_references_alias(h, alias)
{
return true;
}
for (i, j) in joins.iter().enumerate() {
if i == skip_idx {
continue;
}
if let Some(on) = &j.on
&& expr_references_alias(on, alias)
{
return true;
}
}
if let Some(w) = &stmt.where_ {
for p in split_and_conjunctions(w) {
let touches_alias = expr_references_alias(p, alias);
if !touches_alias {
continue;
}
if expr_references_any_other_alias(p, alias) {
return true;
}
}
}
false
}
pub(crate) fn alias_referenced_strictly_elsewhere(
stmt: &SelectStatement,
joins: &[FromJoin],
skip_idx: usize,
alias: &str,
) -> bool {
if alias_referenced_elsewhere(stmt, joins, skip_idx, alias) {
return true;
}
if let Some(w) = &stmt.where_ {
for p in split_and_conjunctions(w) {
if expr_references_alias(p, alias) {
return true;
}
}
}
false
}
pub(crate) fn expr_references_alias(e: &Expr, alias: &str) -> bool {
let mut hit = false;
walk_expr(e, &mut |item| match item {
WalkItem::Column(c) => {
if c.qualifier.as_deref() == Some(alias) {
hit = true;
}
}
WalkItem::Subquery(_) | WalkItem::Unknown => hit = true,
});
hit
}
pub(crate) fn expr_references_any_other_alias(e: &Expr, alias: &str) -> bool {
let mut other = false;
walk_expr(e, &mut |item| match item {
WalkItem::Column(c) => {
if let Some(q) = &c.qualifier
&& q != alias
{
other = true;
}
}
WalkItem::Subquery(_) | WalkItem::Unknown => other = true,
});
other
}
fn build_folded_from(from: &FromClause, drop_indices: &[usize]) -> FromClause {
let drop: alloc::collections::BTreeSet<usize> = drop_indices.iter().copied().collect();
let kept: Vec<FromJoin> = from
.joins
.iter()
.enumerate()
.filter_map(|(i, j)| {
if drop.contains(&i) {
None
} else {
Some(j.clone())
}
})
.collect();
FromClause {
primary: from.primary.clone(),
joins: kept,
}
}
fn rewrite_where(
original: Option<&Expr>,
alias_to_idx: &BTreeMap<String, usize>,
plans: &[FoldPlan],
) -> Option<Expr> {
let folded_idx: alloc::collections::BTreeSet<usize> =
plans.iter().map(|p| p.join_idx + 1).collect();
let mut kept: Vec<Expr> = Vec::new();
if let Some(w) = original {
for p in split_and_conjunctions(w) {
let refs = referenced_aliases(p, alias_to_idx);
if refs.iter().any(|i| folded_idx.contains(i)) {
continue;
}
kept.push(p.clone());
}
}
for p in plans {
if p.pk_values.is_empty() {
kept.push(Expr::Literal(Literal::Bool(false)));
continue;
}
let list: Vec<Expr> = p
.pk_values
.iter()
.map(|v| Expr::Literal(value_to_literal(v)))
.collect();
kept.push(Expr::InList {
expr: alloc::boxed::Box::new(p.outer_expr.clone()),
list,
negated: false,
});
}
if kept.is_empty() {
return None;
}
let mut iter = kept.into_iter();
let mut acc = iter.next().expect("kept non-empty");
for e in iter {
acc = Expr::Binary {
lhs: alloc::boxed::Box::new(acc),
op: BinOp::And,
rhs: alloc::boxed::Box::new(e),
};
}
Some(acc)
}
fn value_to_literal(v: &Value) -> Literal {
match v {
Value::SmallInt(n) => Literal::Integer(i64::from(*n)),
Value::Int(n) => Literal::Integer(i64::from(*n)),
Value::BigInt(n) => Literal::Integer(*n),
Value::Bool(b) => Literal::Bool(*b),
Value::Float(x) => Literal::Float(*x),
Value::Text(s) => Literal::String(s.to_string()),
_ => Literal::Null,
}
}
impl Engine {
fn column_is_single_unique_public(&self, table: &str, col: &str) -> bool {
let Some(t) = self.active_catalog().get(table) else {
return false;
};
let sch = t.schema();
let Some(pos) = sch
.columns
.iter()
.position(|c| c.name.eq_ignore_ascii_case(col))
else {
return false;
};
if sch
.uniqueness_constraints
.iter()
.any(|u| u.columns.as_slice() == [pos])
{
return true;
}
t.index_on(pos).is_some_and(|idx| idx.is_unique)
}
fn collect_unique_keys(
&self,
table_name: &str,
alias: &str,
pk_col: &str,
inner_preds: &[Expr],
cancel: CancelToken<'_>,
) -> Result<Vec<Value<'static>>, EngineError> {
let primary = spg_sql::ast::TableRef {
name: table_name.to_string(),
alias: Some(alias.to_string()),
as_of_segment: None,
unnest_expr: None,
unnest_column_aliases: Vec::new(),
generate_series_args: None,
lateral_subquery: None,
jsonb_each_text_arg: None,
};
let where_ = combine_with_and_owned(inner_preds);
let pk_col_name = ColumnName {
qualifier: Some(alias.to_string()),
name: pk_col.to_string(),
};
let probe = SelectStatement {
ctes: Vec::new(),
distinct: false,
items: alloc::vec![SelectItem::Expr {
expr: Expr::Column(pk_col_name),
alias: None,
}],
from: Some(FromClause {
primary,
joins: Vec::new(),
}),
where_,
group_by: None,
having: None,
unions: Vec::new(),
order_by: Vec::new(),
limit: None,
offset: None,
limit_with_ties: false,
group_by_all: false,
};
let result = self.exec_bare_select_cancel(&probe, cancel)?;
let QueryResult::Rows { rows, .. } = result else {
return Err(EngineError::Unsupported(format!(
"joinfold: inner probe returned non-Rows for {table_name}"
)));
};
let mut out: Vec<Value<'static>> = Vec::with_capacity(rows.len());
for r in rows {
if let Some(v) = r.values.into_iter().next() {
out.push(v);
}
}
Ok(out)
}
}
fn combine_with_and_owned(preds: &[Expr]) -> Option<Expr> {
let mut iter = preds.iter();
let mut acc = iter.next()?.clone();
for e in iter {
acc = Expr::Binary {
lhs: alloc::boxed::Box::new(acc),
op: BinOp::And,
rhs: alloc::boxed::Box::new(e.clone()),
};
}
Some(acc)
}