use alloc::string::String;
use alloc::vec::Vec;
use spg_sql::ast::{
BinOp, ColumnName, Cte, Expr, FromJoin, JoinKind, LimitExpr, Literal, SelectItem,
SelectStatement, TableRef, UnOp,
};
pub static PULLUP_LIMIT1_FIRE_COUNT: core::sync::atomic::AtomicU64 =
core::sync::atomic::AtomicU64::new(0);
pub static BATCHED_SCALAR_KEYED_FIRE_COUNT: core::sync::atomic::AtomicU64 =
core::sync::atomic::AtomicU64::new(0);
pub static BATCHED_SCALAR_KEYED_PROBE_COUNT: core::sync::atomic::AtomicU64 =
core::sync::atomic::AtomicU64::new(0);
pub static BATCHED_SCALAR_FALL_THROUGH_COUNT: core::sync::atomic::AtomicU64 =
core::sync::atomic::AtomicU64::new(0);
pub static EXISTS_PULLUP_CANDIDATE_COUNT: core::sync::atomic::AtomicU64 =
core::sync::atomic::AtomicU64::new(0);
pub static EXISTS_PULLUP_BAIL_INNER_SHAPE: core::sync::atomic::AtomicU64 =
core::sync::atomic::AtomicU64::new(0);
pub static EXISTS_PULLUP_BAIL_INNER_FROM: core::sync::atomic::AtomicU64 =
core::sync::atomic::AtomicU64::new(0);
pub static EXISTS_PULLUP_BAIL_NO_WHERE: core::sync::atomic::AtomicU64 =
core::sync::atomic::AtomicU64::new(0);
pub static EXISTS_PULLUP_BAIL_RESIDUAL_NOT_INNER: core::sync::atomic::AtomicU64 =
core::sync::atomic::AtomicU64::new(0);
pub static EXISTS_PULLUP_BAIL_NO_CORR: core::sync::atomic::AtomicU64 =
core::sync::atomic::AtomicU64::new(0);
pub static EXISTS_PULLUP_BAIL_MULTICOL_DISABLED: core::sync::atomic::AtomicU64 =
core::sync::atomic::AtomicU64::new(0);
pub static EXISTS_PULLUP_BAIL_UNIQUE_KEY_MISSING: core::sync::atomic::AtomicU64 =
core::sync::atomic::AtomicU64::new(0);
pub static EXISTS_PULLUP_FIRE_COUNT: core::sync::atomic::AtomicU64 =
core::sync::atomic::AtomicU64::new(0);
pub static EXISTS_BATCH_FIRE_COUNT: core::sync::atomic::AtomicU64 =
core::sync::atomic::AtomicU64::new(0);
pub static EXISTS_BATCH_FALL_THROUGH_COUNT: core::sync::atomic::AtomicU64 =
core::sync::atomic::AtomicU64::new(0);
pub static EXISTS_PULLUP_MULTICOL_DISABLE: core::sync::atomic::AtomicBool =
core::sync::atomic::AtomicBool::new(false);
use spg_storage::{Row, Value};
use crate::eval::{self, EvalContext};
use crate::substitute::value_to_literal_expr;
use crate::{
CancelToken, Engine, EngineError, QueryResult, aggregate, memoize, order_by_value_cmp, reorder,
value_cmp, visit_expr_columns_and_subqueries,
};
fn scalar_subquery_arity(ncols: usize) -> Result<(), EngineError> {
if ncols == 1 {
Ok(())
} else {
Err(EngineError::Unsupported(
"subquery must return only one column".into(),
))
}
}
fn build_row_comparison(row: &[Expr], op: spg_sql::ast::BinOp, rhs: &[Expr]) -> Expr {
use alloc::boxed::Box;
use spg_sql::ast::{BinOp, UnOp};
fn row_eq(lhs: &[Expr], rhs: &[Expr]) -> Expr {
let mut it = lhs.iter().zip(rhs.iter()).map(|(l, r)| Expr::Binary {
lhs: Box::new(l.clone()),
op: BinOp::Eq,
rhs: Box::new(r.clone()),
});
let first = it.next().expect("row has >= 1 element");
it.fold(first, |acc, e| Expr::Binary {
lhs: Box::new(acc),
op: BinOp::And,
rhs: Box::new(e),
})
}
fn row_lex(lhs: &[Expr], rhs: &[Expr], strict: BinOp, last: BinOp) -> Expr {
if lhs.len() == 1 {
return Expr::Binary {
lhs: Box::new(lhs[0].clone()),
op: last,
rhs: Box::new(rhs[0].clone()),
};
}
let head_strict = Expr::Binary {
lhs: Box::new(lhs[0].clone()),
op: strict,
rhs: Box::new(rhs[0].clone()),
};
let head_eq = Expr::Binary {
lhs: Box::new(lhs[0].clone()),
op: BinOp::Eq,
rhs: Box::new(rhs[0].clone()),
};
Expr::Binary {
lhs: Box::new(head_strict),
op: BinOp::Or,
rhs: Box::new(Expr::Binary {
lhs: Box::new(head_eq),
op: BinOp::And,
rhs: Box::new(row_lex(&lhs[1..], &rhs[1..], strict, last)),
}),
}
}
match op {
BinOp::Eq => row_eq(row, rhs),
BinOp::NotEq => Expr::Unary {
op: UnOp::Not,
expr: Box::new(row_eq(row, rhs)),
},
BinOp::Lt => row_lex(row, rhs, BinOp::Lt, BinOp::Lt),
BinOp::LtEq => row_lex(row, rhs, BinOp::Lt, BinOp::LtEq),
BinOp::Gt => row_lex(row, rhs, BinOp::Gt, BinOp::Gt),
BinOp::GtEq => row_lex(row, rhs, BinOp::Gt, BinOp::GtEq),
_ => Expr::Literal(Literal::Bool(false)), }
}
impl Engine {
pub(crate) fn eval_expr_with_correlated(
&self,
expr: &Expr,
row: &Row<'static>,
ctx: &EvalContext<'_>,
cancel: CancelToken<'_>,
mut memo: Option<&mut memoize::MemoizeCache>,
) -> Result<Value<'static>, EngineError> {
let has_subq = if let Some(m) = memo.as_deref_mut() {
let key = core::ptr::from_ref::<Expr>(expr) as usize;
match m.has_subquery.get(&key) {
Some(b) => *b,
None => {
let b = expr_has_subquery(expr);
m.has_subquery.insert(key, b);
b
}
}
} else {
expr_has_subquery(expr)
};
if !has_subq {
if let Some(m) = memo.as_deref_mut()
&& expr_may_use_in_set(expr)
{
return eval_with_in_sets(expr, row, ctx, m);
}
return eval::eval_expr(expr, row, ctx).map_err(EngineError::Eval);
}
if let Some(m) = memo.as_deref_mut() {
let key = core::ptr::from_ref::<Expr>(expr) as usize;
let plan_hit = m.expr_plans.contains_key(&key);
let exists_plan_hit = m.exists_plans.contains_key(&key);
let mut subs: Vec<&SelectStatement> = Vec::new();
let mut exists_subs: Vec<&SelectStatement> = Vec::new();
if !plan_hit {
collect_scalar_subqueries(expr, &mut subs);
}
if !exists_plan_hit {
collect_exists_subqueries(expr, &mut exists_subs);
}
if !plan_hit && !subs.is_empty() {
let mut plan: Vec<Option<alloc::rc::Rc<memoize::GroupMap>>> =
Vec::with_capacity(subs.len());
for sub in &subs {
let repr = alloc::format!("{sub}");
if !m.group_maps.contains_key(&repr) {
let built = self
.try_batch_correlated_scalar(sub, None, cancel)?
.map(alloc::rc::Rc::new);
m.group_maps.insert(repr.clone(), built);
}
plan.push(m.group_maps.get(&repr).cloned().flatten());
}
let mut template = expr.clone();
hollow_scalar_subqueries(&mut template);
m.expr_plans.insert(key, (subs.len(), plan, template));
}
if !exists_plan_hit && !exists_subs.is_empty() {
let mut eplan: Vec<Option<alloc::rc::Rc<memoize::ExistsSet>>> =
Vec::with_capacity(exists_subs.len());
for sub in &exists_subs {
let built = self
.try_batch_correlated_exists(sub, cancel)?
.map(alloc::rc::Rc::new);
if built.is_some() {
EXISTS_BATCH_FIRE_COUNT.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
} else {
EXISTS_BATCH_FALL_THROUGH_COUNT
.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
}
eplan.push(built);
}
m.exists_plans.insert(key, eplan);
}
if !m.expr_plans.contains_key(&key)
&& let Some((negated, wrapped_in_not)) = bare_exists_shape(expr)
&& let Some(plan) = m.exists_plans.get(&key)
&& plan.len() == 1
&& let Some(Some(es)) = plan.first()
{
let bit = planned_exists_bit(es, negated, row, ctx)?;
return Ok(Value::Bool(if wrapped_in_not { !bit } else { bit }));
}
let scalar_ready = m
.expr_plans
.get(&key)
.map(|(_, plan, _)| !plan.is_empty() && plan.iter().all(|p| p.is_some()))
.unwrap_or(false);
let exists_ready = m
.exists_plans
.get(&key)
.map(|plan| !plan.is_empty() && plan.iter().all(|p| p.is_some()))
.unwrap_or(false);
if scalar_ready || exists_ready {
let scalar_plan = m
.expr_plans
.get(&key)
.map(|(_, plan, template)| (plan.clone(), template.clone()));
let exists_plan = m.exists_plans.get(&key).cloned();
let mut e = match &scalar_plan {
Some((_, template)) => template.clone(),
None => expr.clone(),
};
let mut all_ok = true;
if let Some((plan, _)) = &scalar_plan {
let mut idx = 0usize;
all_ok &= splice_planned_subqueries(&mut e, plan, &mut idx, row, ctx)?;
}
if all_ok && let Some(plan) = &exists_plan {
let mut idx = 0usize;
all_ok &= splice_planned_exists(&mut e, plan, &mut idx, row, ctx)?;
}
if all_ok {
if expr_has_subquery(&e) {
self.resolve_correlated_in_expr(&mut e, row, ctx, cancel, memo)?;
}
return eval::eval_expr(&e, row, ctx).map_err(EngineError::Eval);
}
}
}
let mut e = expr.clone();
self.resolve_correlated_in_expr(&mut e, row, ctx, cancel, memo)?;
eval::eval_expr(&e, row, ctx).map_err(EngineError::Eval)
}
pub(crate) fn materialize_quantified_rows(
&self,
inner: &SelectStatement,
cancel: CancelToken<'_>,
) -> Result<Expr, EngineError> {
let r = self.exec_select_cancel(inner, cancel)?;
let QueryResult::Rows { rows, .. } = r else {
return Err(EngineError::Unsupported(
"ANY/ALL subquery: inner did not return rows".into(),
));
};
let mut items = alloc::vec::Vec::with_capacity(rows.len());
for r0 in rows {
let v = r0.values.into_iter().next().unwrap_or(Value::Null);
items.push(value_to_literal_expr(v)?);
}
Ok(Expr::Array(items))
}
fn resolve_correlated_in_expr(
&self,
e: &mut Expr,
row: &Row<'static>,
ctx: &EvalContext<'_>,
cancel: CancelToken<'_>,
mut memo: Option<&mut memoize::MemoizeCache>,
) -> Result<(), EngineError> {
match e {
Expr::NamedArg { expr, .. } | Expr::Variadic(expr) => {
self.resolve_correlated_in_expr(expr, row, ctx, cancel, memo.as_deref_mut())?;
}
Expr::AggregateOrdered { call, order_by, .. } => {
self.resolve_correlated_in_expr(call, row, ctx, cancel, memo.as_deref_mut())?;
for o in order_by.iter_mut() {
self.resolve_correlated_in_expr(
&mut o.expr,
row,
ctx,
cancel,
memo.as_deref_mut(),
)?;
}
}
Expr::ScalarSubquery(inner) => {
if memo.is_some() {
let ptr_key = core::ptr::from_ref::<SelectStatement>(&**inner) as usize;
let entry_known = memo
.as_ref()
.is_some_and(|m| m.group_maps_by_ptr.contains_key(&ptr_key));
if !entry_known {
let built = self
.try_batch_correlated_scalar(inner, None, cancel)?
.map(alloc::rc::Rc::new);
if let Some(m) = memo.as_deref_mut() {
m.group_maps_by_ptr.insert(ptr_key, built);
}
}
if let Some(m) = memo.as_deref_mut()
&& let Some(Some(gm)) = m.group_maps_by_ptr.get(&ptr_key)
{
let (outer_col, map, empty_default) = gm.as_ref();
let key_v = eval::eval_expr(&Expr::Column(outer_col.clone()), row, ctx)
.map_err(EngineError::Eval)?;
let v = map
.get(&aggregate::encode_key(core::slice::from_ref(&key_v)))
.cloned()
.unwrap_or_else(|| empty_default.clone());
*e = value_to_literal_expr(v)?;
return Ok(());
}
}
let cache_key = memo.as_ref().map(|_| memoize::CacheKey {
subquery_repr: alloc::format!("{}", **inner),
outer_values: row.values.iter().cloned().map(Value::into_owned).collect(),
});
if let (Some(cache), Some(k)) = (memo.as_deref_mut(), cache_key.as_ref())
&& let Some(cached) = cache.get(k)
{
*e = value_to_literal_expr(cached)?;
return Ok(());
}
if let Some(v) = self.try_scalar_count_pk_eq_probe(inner, row, ctx)? {
*e = value_to_literal_expr(v)?;
return Ok(());
}
let mut s = (**inner).clone();
substitute_outer_columns(&mut s, row, ctx, self.active_catalog());
let r = self.exec_select_cancel(&s, cancel)?;
let QueryResult::Rows { columns, rows, .. } = r else {
return Err(EngineError::Unsupported(
"scalar subquery: inner did not return rows".into(),
));
};
scalar_subquery_arity(columns.len())?;
let value = match rows.as_slice() {
[] => Value::Null,
[r0] => r0.values.first().cloned().unwrap_or(Value::Null),
_ => {
return Err(EngineError::CardinalityViolation);
}
};
if let (Some(cache), Some(k)) = (memo.as_deref_mut(), cache_key) {
cache.insert(k, value.clone());
}
*e = value_to_literal_expr(value)?;
}
Expr::Exists { subquery, negated } => {
if memo.is_some() {
let repr = alloc::format!("{}", **subquery);
let known = memo
.as_ref()
.is_some_and(|m| m.exists_sets.contains_key(&repr));
if !known {
let built = self
.try_batch_correlated_exists(subquery, cancel)?
.map(alloc::rc::Rc::new);
if let Some(m) = memo.as_deref_mut() {
m.exists_sets.insert(repr.clone(), built);
}
}
if let Some(m) = memo.as_deref_mut()
&& let Some(Some(es)) = m.exists_sets.get(&repr)
{
let (outer_cols, set) = es.as_ref();
let mut key_vals: Vec<Value<'static>> =
Vec::with_capacity(outer_cols.len());
let mut any_null = false;
for oc in outer_cols {
let v = eval::eval_expr(oc, row, ctx).map_err(EngineError::Eval)?;
if matches!(v, Value::Null) {
any_null = true;
}
key_vals.push(v);
}
let present =
!any_null && set.contains(&aggregate::encode_canonical_key(&key_vals));
let bit = if *negated { !present } else { present };
*e = Expr::Literal(Literal::Bool(bit));
return Ok(());
}
}
let mut s = (**subquery).clone();
substitute_outer_columns(&mut s, row, ctx, self.active_catalog());
let r = self.exec_select_cancel(&s, cancel)?;
let exists = matches!(r, QueryResult::Rows { rows, .. } if !rows.is_empty());
let bit = if *negated { !exists } else { exists };
*e = Expr::Literal(Literal::Bool(bit));
}
Expr::InSubquery {
expr: lhs,
subquery,
negated,
} => {
self.resolve_correlated_in_expr(lhs, row, ctx, cancel, memo.as_deref_mut())?;
let lhs_val = eval::eval_expr(lhs, row, ctx).map_err(EngineError::Eval)?;
let mut s = (**subquery).clone();
substitute_outer_columns(&mut s, row, ctx, self.active_catalog());
let r = self.exec_select_cancel(&s, cancel)?;
let QueryResult::Rows { columns, rows, .. } = r else {
return Err(EngineError::Unsupported(
"IN-subquery: inner did not return rows".into(),
));
};
if columns.len() != 1 {
return Err(EngineError::Unsupported(
if columns.is_empty() {
"subquery has too few columns"
} else {
"subquery has too many columns"
}
.into(),
));
}
let mut found = false;
let mut any_null = false;
for r0 in rows {
let v = r0.values.into_iter().next().unwrap_or(Value::Null);
if v.is_null() {
any_null = true;
continue;
}
if value_cmp(&v, &lhs_val) == core::cmp::Ordering::Equal {
found = true;
break;
}
}
if !found && any_null {
*e = Expr::Literal(Literal::Null);
return Ok(());
}
let bit = if found { !*negated } else { *negated };
*e = Expr::Literal(Literal::Bool(bit));
}
Expr::RowInSubquery {
row: row_exprs,
subquery,
negated,
} => {
for el in row_exprs.iter_mut() {
self.resolve_correlated_in_expr(el, row, ctx, cancel, memo.as_deref_mut())?;
}
let lhs_vals: Vec<Value> = row_exprs
.iter()
.map(|el| eval::eval_expr(el, row, ctx).map_err(EngineError::Eval))
.collect::<Result<_, _>>()?;
let mut s = (**subquery).clone();
substitute_outer_columns(&mut s, row, ctx, self.active_catalog());
let r = self.exec_select_cancel(&s, cancel)?;
let QueryResult::Rows { columns, rows, .. } = r else {
return Err(EngineError::Unsupported(
"row IN-subquery: inner did not return rows".into(),
));
};
if columns.len() != lhs_vals.len() {
return Err(EngineError::Unsupported(alloc::format!(
"row IN-subquery: left side has {} column(s), subquery returns {}",
lhs_vals.len(),
columns.len()
)));
}
let mut found = false;
let mut any_null = false;
'rows: for r0 in rows {
let mut has_null = false;
for (j, sub_v) in r0.values.iter().enumerate() {
let lv = &lhs_vals[j];
if lv.is_null() || sub_v.is_null() {
has_null = true;
} else if value_cmp(lv, sub_v) != core::cmp::Ordering::Equal {
continue 'rows; }
}
if has_null {
any_null = true; } else {
found = true; break;
}
}
if !found && any_null {
*e = Expr::Literal(Literal::Null);
return Ok(());
}
let bit = if found { !*negated } else { *negated };
*e = Expr::Literal(Literal::Bool(bit));
}
Expr::RowCmpSubquery {
row: row_exprs,
op,
subquery,
} => {
for el in row_exprs.iter_mut() {
self.resolve_correlated_in_expr(el, row, ctx, cancel, memo.as_deref_mut())?;
}
let mut s = (**subquery).clone();
substitute_outer_columns(&mut s, row, ctx, self.active_catalog());
let r = self.exec_select_cancel(&s, cancel)?;
let QueryResult::Rows {
columns, mut rows, ..
} = r
else {
return Err(EngineError::Unsupported(
"row comparison subquery: inner did not return rows".into(),
));
};
if rows.is_empty() {
*e = Expr::Literal(Literal::Null);
return Ok(());
}
if rows.len() > 1 {
return Err(EngineError::CardinalityViolation);
}
if columns.len() != row_exprs.len() {
return Err(EngineError::Unsupported(alloc::format!(
"row comparison: left side has {} column(s), subquery returns {}",
row_exprs.len(),
columns.len()
)));
}
let rhs: Vec<Expr> = rows
.remove(0)
.values
.into_iter()
.map(value_to_literal_expr)
.collect::<Result<_, _>>()?;
let cmp = build_row_comparison(row_exprs, *op, &rhs);
let v = eval::eval_expr(&cmp, row, ctx).map_err(EngineError::Eval)?;
*e = value_to_literal_expr(v)?;
}
Expr::Binary { lhs, rhs, .. } => {
self.resolve_correlated_in_expr(lhs, row, ctx, cancel, memo.as_deref_mut())?;
self.resolve_correlated_in_expr(rhs, row, ctx, cancel, memo.as_deref_mut())?;
}
Expr::Unary { expr, .. }
| Expr::Cast { expr, .. }
| Expr::IsNull { expr, .. }
| Expr::BoolTest { expr, .. }
| Expr::FieldAccess { base: expr, .. } => {
self.resolve_correlated_in_expr(expr, row, ctx, cancel, memo.as_deref_mut())?;
}
Expr::Like { expr, pattern, .. } => {
self.resolve_correlated_in_expr(expr, row, ctx, cancel, memo.as_deref_mut())?;
self.resolve_correlated_in_expr(pattern, row, ctx, cancel, memo.as_deref_mut())?;
}
Expr::FunctionCall { args, .. } => {
for a in args {
self.resolve_correlated_in_expr(a, row, ctx, cancel, memo.as_deref_mut())?;
}
}
Expr::Extract { source, .. } => {
self.resolve_correlated_in_expr(source, row, ctx, cancel, memo.as_deref_mut())?;
}
Expr::WindowFunction { .. }
| Expr::Literal(_)
| Expr::Placeholder(_)
| Expr::Column(_) => {}
Expr::Array(items) => {
for elem in items {
self.resolve_correlated_in_expr(elem, row, ctx, cancel, memo.as_deref_mut())?;
}
}
Expr::ArraySubscript { target, index } => {
self.resolve_correlated_in_expr(target, row, ctx, cancel, memo.as_deref_mut())?;
self.resolve_correlated_in_expr(index, row, ctx, cancel, memo.as_deref_mut())?;
}
Expr::ArraySlice { target, lo, hi } => {
self.resolve_correlated_in_expr(target, row, ctx, cancel, memo.as_deref_mut())?;
if let Some(l) = lo {
self.resolve_correlated_in_expr(l, row, ctx, cancel, memo.as_deref_mut())?;
}
if let Some(h) = hi {
self.resolve_correlated_in_expr(h, row, ctx, cancel, memo.as_deref_mut())?;
}
}
Expr::AnyAll { expr, array, .. } => {
self.resolve_correlated_in_expr(expr, row, ctx, cancel, memo.as_deref_mut())?;
if let Expr::ScalarSubquery(inner) = array.as_mut() {
let mut s = (**inner).clone();
substitute_outer_columns(&mut s, row, ctx, self.active_catalog());
**array = self.materialize_quantified_rows(&s, cancel)?;
} else {
self.resolve_correlated_in_expr(array, row, ctx, cancel, memo.as_deref_mut())?;
}
}
Expr::InList { expr, list, .. } => {
self.resolve_correlated_in_expr(expr, row, ctx, cancel, memo.as_deref_mut())?;
for item in list {
self.resolve_correlated_in_expr(item, row, ctx, cancel, memo.as_deref_mut())?;
}
}
Expr::Case {
operand,
branches,
else_branch,
} => {
if let Some(o) = operand {
self.resolve_correlated_in_expr(o, row, ctx, cancel, memo.as_deref_mut())?;
}
for (w, t) in branches {
self.resolve_correlated_in_expr(w, row, ctx, cancel, memo.as_deref_mut())?;
self.resolve_correlated_in_expr(t, row, ctx, cancel, memo.as_deref_mut())?;
}
if let Some(e) = else_branch {
self.resolve_correlated_in_expr(e, row, ctx, cancel, memo.as_deref_mut())?;
}
}
}
Ok(())
}
#[inline(never)]
fn arm_row_cmp_subquery(
&self,
e: &Expr,
cancel: CancelToken<'_>,
) -> Result<Option<Expr>, EngineError> {
let Expr::RowCmpSubquery { row, op, subquery } = e else {
return Ok(None);
};
if select_is_correlated(subquery) {
return Ok(None);
}
let mut s = (**subquery).clone();
self.resolve_select_subqueries(&mut s, cancel)?;
let r = match self.exec_select_cancel(&s, cancel) {
Ok(r) => r,
Err(e) if is_correlation_error(&e) => return Ok(None),
Err(e) => return Err(e),
};
let QueryResult::Rows {
columns, mut rows, ..
} = r
else {
return Err(EngineError::Unsupported(
"row comparison subquery: inner statement did not return rows".into(),
));
};
if rows.is_empty() {
return Ok(Some(Expr::Literal(Literal::Null)));
}
if rows.len() > 1 {
return Err(EngineError::CardinalityViolation);
}
if columns.len() != row.len() {
return Err(EngineError::Unsupported(alloc::format!(
"row comparison: left side has {} column(s), subquery returns {}",
row.len(),
columns.len()
)));
}
let rhs: Vec<Expr> = rows
.remove(0)
.values
.into_iter()
.map(value_to_literal_expr)
.collect::<Result<_, _>>()?;
Ok(Some(build_row_comparison(row, *op, &rhs)))
}
#[inline(never)]
fn arm_row_in_subquery(
&self,
e: &Expr,
cancel: CancelToken<'_>,
) -> Result<Option<Expr>, EngineError> {
let Expr::RowInSubquery {
row,
subquery,
negated,
} = e
else {
return Ok(None);
};
use alloc::boxed::Box;
if select_is_correlated(subquery) {
return Ok(None);
}
let mut s = (**subquery).clone();
self.resolve_select_subqueries(&mut s, cancel)?;
let r = match self.exec_select_cancel(&s, cancel) {
Ok(r) => r,
Err(e) if is_correlation_error(&e) => return Ok(None),
Err(e) => return Err(e),
};
let QueryResult::Rows { columns, rows, .. } = r else {
return Err(EngineError::Unsupported(
"row IN-subquery: inner statement did not return rows".into(),
));
};
if columns.len() != row.len() {
return Err(EngineError::Unsupported(alloc::format!(
"row IN-subquery: left side has {} column(s), subquery returns {}",
row.len(),
columns.len()
)));
}
let mut alts: Vec<Expr> = Vec::with_capacity(rows.len());
for r0 in rows {
let mut conj: Option<Expr> = None;
for (lhs_el, v) in row.iter().zip(r0.values) {
let eq = Expr::Binary {
lhs: Box::new(lhs_el.clone()),
op: BinOp::Eq,
rhs: Box::new(value_to_literal_expr(v)?),
};
conj = Some(match conj {
None => eq,
Some(prev) => Expr::Binary {
lhs: Box::new(prev),
op: BinOp::And,
rhs: Box::new(eq),
},
});
}
if let Some(c) = conj {
alts.push(c);
}
}
let combined = match alts.into_iter().reduce(|acc, e| Expr::Binary {
lhs: Box::new(acc),
op: BinOp::Or,
rhs: Box::new(e),
}) {
Some(c) => c,
None => Expr::Literal(Literal::Bool(false)),
};
let result = if *negated {
Expr::Unary {
op: UnOp::Not,
expr: Box::new(combined),
}
} else {
combined
};
Ok(Some(result))
}
#[inline(never)]
fn arm_in_subquery(
&self,
e: &Expr,
cancel: CancelToken<'_>,
) -> Result<Option<Expr>, EngineError> {
let Expr::InSubquery {
expr,
subquery,
negated,
} = e
else {
return Ok(None);
};
if select_is_correlated(subquery) {
return Ok(None);
}
let mut s = (**subquery).clone();
self.resolve_select_subqueries(&mut s, cancel)?;
let r = match self.exec_select_cancel(&s, cancel) {
Ok(r) => r,
Err(e) if is_correlation_error(&e) => return Ok(None),
Err(e) => return Err(e),
};
let QueryResult::Rows { columns, rows, .. } = r else {
return Err(EngineError::Unsupported(
"IN-subquery: inner statement did not return rows".into(),
));
};
if columns.len() != 1 {
return Err(EngineError::Unsupported(
if columns.is_empty() {
"subquery has too few columns"
} else {
"subquery has too many columns"
}
.into(),
));
}
let mut list: Vec<Expr> = Vec::with_capacity(rows.len());
for row in rows {
let v = row.values.into_iter().next().unwrap_or(Value::Null);
list.push(value_to_literal_expr(v)?);
}
Ok(Some(Expr::InList {
expr: expr.clone(),
list,
negated: *negated,
}))
}
#[inline(never)]
fn arm_exists(&self, e: &Expr, cancel: CancelToken<'_>) -> Result<Option<Expr>, EngineError> {
let Expr::Exists { subquery, negated } = e else {
return Ok(None);
};
if select_is_correlated(subquery) {
return Ok(None);
}
let mut s = (**subquery).clone();
self.resolve_select_subqueries(&mut s, cancel)?;
let r = match self.exec_select_cancel(&s, cancel) {
Ok(r) => r,
Err(e) if is_correlation_error(&e) => return Ok(None),
Err(e) => return Err(e),
};
let exists = match r {
QueryResult::Rows { rows, .. } => !rows.is_empty(),
QueryResult::CommandOk { .. } => false,
};
let bit = if *negated { !exists } else { exists };
Ok(Some(Expr::Literal(Literal::Bool(bit))))
}
#[inline(never)]
fn arm_scalar_subquery(
&self,
e: &Expr,
cancel: CancelToken<'_>,
) -> Result<Option<Expr>, EngineError> {
let Expr::ScalarSubquery(inner) = e else {
return Ok(None);
};
if select_is_correlated(inner) {
return Ok(None);
}
let mut s = (**inner).clone();
self.resolve_select_subqueries(&mut s, cancel)?;
let r = match self.exec_select_cancel(&s, cancel) {
Ok(r) => r,
Err(e) if is_correlation_error(&e) => return Ok(None),
Err(e) => return Err(e),
};
let QueryResult::Rows { columns, rows, .. } = r else {
return Err(EngineError::Unsupported(
"scalar subquery: inner statement did not return rows".into(),
));
};
scalar_subquery_arity(columns.len())?;
let value = match rows.as_slice() {
[] => Value::Null,
[row] => row.values.first().cloned().unwrap_or(Value::Null),
_ => {
return Err(EngineError::CardinalityViolation);
}
};
Ok(Some(value_to_literal_expr(value)?))
}
pub(crate) fn subquery_replacement(
&self,
e: &Expr,
cancel: CancelToken<'_>,
) -> Result<Option<Expr>, EngineError> {
match e {
Expr::ScalarSubquery(..) => self.arm_scalar_subquery(e, cancel),
Expr::Exists { .. } => self.arm_exists(e, cancel),
Expr::InSubquery { .. } => self.arm_in_subquery(e, cancel),
Expr::RowInSubquery { .. } => self.arm_row_in_subquery(e, cancel),
Expr::RowCmpSubquery { .. } => self.arm_row_cmp_subquery(e, cancel),
_ => Ok(None),
}
}
}
impl Engine {
pub(crate) fn try_batch_correlated_scalar(
&self,
inner: &SelectStatement,
restrict: Option<(&[Row<'static>], &EvalContext<'_>)>,
cancel: CancelToken<'_>,
) -> Result<Option<memoize::GroupMap>, EngineError> {
use spg_sql::ast::{BinOp, SelectItem as SI};
if !inner.ctes.is_empty()
|| !inner.unions.is_empty()
|| inner.group_by.is_some()
|| inner.having.is_some()
|| inner.distinct
|| inner.items.len() != 1
|| inner.order_by.len() > 1
|| inner.offset.is_some()
{
return Ok(None);
}
if let Some(le) = &inner.limit
&& le.as_literal() != Some(1)
{
return Ok(None);
}
let Some(from) = &inner.from else {
return Ok(None);
};
if from.primary.lateral_subquery.is_some() || from.primary.unnest_expr.is_some() {
return Ok(None);
}
let mut inner_aliases: Vec<String> = Vec::new();
inner_aliases.push(
from.primary
.alias
.clone()
.unwrap_or_else(|| from.primary.name.clone()),
);
for j in &from.joins {
if j.table.lateral_subquery.is_some() || j.table.unnest_expr.is_some() {
return Ok(None);
}
inner_aliases.push(
j.table
.alias
.clone()
.unwrap_or_else(|| j.table.name.clone()),
);
}
let is_inner = |c: &spg_sql::ast::ColumnName| -> bool {
match &c.qualifier {
Some(q) => inner_aliases.iter().any(|a| a.eq_ignore_ascii_case(q)),
None => false,
}
};
let is_outer = |c: &spg_sql::ast::ColumnName| -> bool {
match &c.qualifier {
Some(q) => !inner_aliases.iter().any(|a| a.eq_ignore_ascii_case(q)),
None => c.name.starts_with("__grp_") || c.name.starts_with("__agg_"),
}
};
let all_inner = |e: &Expr| -> bool {
let mut cols: Vec<spg_sql::ast::ColumnName> = Vec::new();
let mut subs: Vec<&SelectStatement> = Vec::new();
visit_expr_columns_and_subqueries(e, &mut |c| cols.push(c.clone()), &mut |sub| {
subs.push(sub)
});
subs.is_empty() && cols.iter().all(|c| is_inner(c) && !c.name.is_empty())
};
let Some(w) = &inner.where_ else {
return Ok(None);
};
let conjuncts = reorder::split_and_conjunctions(w);
let mut corr: Option<(spg_sql::ast::ColumnName, spg_sql::ast::ColumnName)> = None; let mut rest: Vec<&Expr> = Vec::new();
for c in conjuncts {
if let Expr::Binary {
lhs,
op: BinOp::Eq,
rhs,
} = c
&& let (Expr::Column(a), Expr::Column(b)) = (lhs.as_ref(), rhs.as_ref())
{
let pair = if is_inner(a) && is_outer(b) {
Some((a.clone(), b.clone()))
} else if is_inner(b) && is_outer(a) {
Some((b.clone(), a.clone()))
} else {
None
};
if let Some(p) = pair {
if corr.is_some() {
return Ok(None); }
corr = Some(p);
continue;
}
}
if !all_inner(c) {
return Ok(None);
}
rest.push(c);
}
let Some((inner_col, outer_col)) = corr else {
return Ok(None);
};
let SI::Expr { expr: out_expr, .. } = &inner.items[0] else {
return Ok(None);
};
if !all_inner(out_expr) {
return Ok(None);
}
let order = inner.order_by.first();
if let Some(o) = order
&& !all_inner(&o.expr)
{
return Ok(None);
}
let mut batch = inner.clone();
batch.limit = None;
batch.offset = None;
batch.order_by = Vec::new();
batch.where_ = rest
.iter()
.map(|e| (*e).clone())
.reduce(|a, b| Expr::Binary {
lhs: alloc::boxed::Box::new(a),
op: BinOp::And,
rhs: alloc::boxed::Box::new(b),
});
let mut items: Vec<SI> = alloc::vec![SI::Expr {
expr: Expr::Column(inner_col.clone()),
alias: None,
}];
if let Some(o) = order {
items.push(SI::Expr {
expr: o.expr.clone(),
alias: None,
});
}
items.push(SI::Expr {
expr: out_expr.clone(),
alias: None,
});
batch.items = items;
if aggregate::contains_aggregate(out_expr) {
batch.group_by = Some(alloc::vec![Expr::Column(inner_col.clone())]);
}
let keyed: Option<(&[Row<'static>], &EvalContext<'_>)> =
restrict.and_then(|(rows, rctx)| {
let driver_name: &str = if from.joins.is_empty() {
from.primary.name.as_str()
} else {
let q = inner_col.qualifier.as_deref()?;
let primary_alias = from
.primary
.alias
.as_deref()
.unwrap_or(from.primary.name.as_str());
if primary_alias.eq_ignore_ascii_case(q) {
from.primary.name.as_str()
} else {
from.joins
.iter()
.find(|j| {
j.table
.alias
.as_deref()
.unwrap_or(j.table.name.as_str())
.eq_ignore_ascii_case(q)
})
.map(|j| j.table.name.as_str())?
}
};
let table = self.active_catalog().get(driver_name)?;
let pos = table
.schema()
.columns
.iter()
.position(|c| c.name.eq_ignore_ascii_case(&inner_col.name))?;
table.index_on(pos)?;
if rows.len().saturating_mul(4) >= table.row_count() {
return None;
}
if !from.joins.is_empty() {
let driver_alias = inner_col.qualifier.as_deref()?;
if !reorder::drive_from(&mut batch, driver_alias) {
return None;
}
}
Some((rows, rctx))
});
let rows = if let Some((restrict_rows, rctx)) = keyed {
BATCHED_SCALAR_KEYED_FIRE_COUNT.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
let mut seen: alloc::collections::BTreeSet<String> =
alloc::collections::BTreeSet::new();
let mut key_lits: Vec<Expr> = Vec::new();
for srow in restrict_rows {
cancel.check()?;
let kv = eval::eval_expr(&Expr::Column(outer_col.clone()), srow, rctx)
.map_err(EngineError::Eval)?;
if matches!(kv, Value::Null) {
continue;
}
if !seen.insert(aggregate::encode_key(core::slice::from_ref(&kv))) {
continue;
}
key_lits.push(value_to_literal_expr(kv)?);
}
if key_lits.is_empty() {
Vec::new()
} else {
let in_pred = Expr::InList {
expr: alloc::boxed::Box::new(Expr::Column(inner_col.clone())),
list: key_lits,
negated: false,
};
let mut probe = batch.clone();
probe.where_ = Some(match probe.where_.take() {
Some(w) => Expr::Binary {
lhs: alloc::boxed::Box::new(w),
op: BinOp::And,
rhs: alloc::boxed::Box::new(in_pred),
},
None => in_pred,
});
BATCHED_SCALAR_KEYED_PROBE_COUNT
.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
if let QueryResult::Rows { rows, .. } = self.exec_select_cancel(&probe, cancel)? {
rows
} else {
Vec::new()
}
}
} else {
BATCHED_SCALAR_FALL_THROUGH_COUNT.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
let r = self.exec_select_cancel(&batch, cancel)?;
let QueryResult::Rows { rows, .. } = r else {
return Ok(None);
};
rows
};
let has_order = order.is_some();
let (desc, nf) = order
.map(|o| (o.desc, o.nulls_first))
.unwrap_or((false, None));
let mut best: alloc::collections::BTreeMap<String, (Option<Value>, Value)> =
alloc::collections::BTreeMap::new();
for row in rows {
let key_v = row.values.first().cloned().unwrap_or(Value::Null);
if matches!(key_v, Value::Null) {
continue;
}
let key = aggregate::encode_key(core::slice::from_ref(&key_v));
let (ord_v, out_v) = if has_order {
(
Some(row.values.get(1).cloned().unwrap_or(Value::Null)),
row.values.get(2).cloned().unwrap_or(Value::Null),
)
} else {
(None, row.values.get(1).cloned().unwrap_or(Value::Null))
};
match best.get(&key) {
None => {
best.insert(key, (ord_v, out_v));
}
Some((cur_ord, _)) if has_order => {
let cand = ord_v.clone().unwrap_or(Value::Null);
let cur = cur_ord.clone().unwrap_or(Value::Null);
if order_by_value_cmp(desc, nf, &cand, &cur) == core::cmp::Ordering::Less {
best.insert(key, (ord_v, out_v));
}
}
Some(_) => {} }
}
let map = best.into_iter().map(|(k, (_, v))| (k, v)).collect();
let empty_default = scalar_subquery_empty_default(inner);
Ok(Some((outer_col, map, empty_default)))
}
}
impl Engine {
pub(crate) fn try_batch_correlated_exists(
&self,
inner: &SelectStatement,
cancel: CancelToken<'_>,
) -> Result<Option<memoize::ExistsSet>, EngineError> {
use spg_sql::ast::SelectItem as SI;
if !inner.ctes.is_empty()
|| !inner.unions.is_empty()
|| inner.group_by.is_some()
|| inner.having.is_some()
|| inner.distinct
{
return Ok(None);
}
let Some(from) = &inner.from else {
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 mut inner_aliases: Vec<String> = Vec::new();
inner_aliases.push(
from.primary
.alias
.clone()
.unwrap_or_else(|| from.primary.name.clone()),
);
for j in &from.joins {
if j.table.lateral_subquery.is_some() || j.table.unnest_expr.is_some() {
return Ok(None);
}
inner_aliases.push(
j.table
.alias
.clone()
.unwrap_or_else(|| j.table.name.clone()),
);
}
let is_inner = |c: &spg_sql::ast::ColumnName| -> bool {
match &c.qualifier {
Some(q) => inner_aliases.iter().any(|a| a.eq_ignore_ascii_case(q)),
None => false,
}
};
let is_outer = |c: &spg_sql::ast::ColumnName| -> bool {
match &c.qualifier {
Some(q) => !inner_aliases.iter().any(|a| a.eq_ignore_ascii_case(q)),
None => c.name.starts_with("__grp_") || c.name.starts_with("__agg_"),
}
};
let all_inner = |e: &Expr| -> bool {
let mut cols: Vec<spg_sql::ast::ColumnName> = Vec::new();
let mut subs: Vec<&SelectStatement> = Vec::new();
visit_expr_columns_and_subqueries(e, &mut |c| cols.push(c.clone()), &mut |sub| {
subs.push(sub)
});
subs.is_empty() && cols.iter().all(|c| is_inner(c) && !c.name.is_empty())
};
let Some(w) = &inner.where_ else {
return Ok(None);
};
let conjuncts = reorder::split_and_conjunctions(w);
let outer_only_key = |e: &Expr| -> bool {
fn shape(e: &Expr, is_outer: &dyn Fn(&spg_sql::ast::ColumnName) -> bool) -> bool {
use spg_sql::ast::BinOp as B;
match e {
Expr::Column(c) => is_outer(c),
Expr::Literal(_) => true,
Expr::Unary { expr, .. } | Expr::Cast { expr, .. } => shape(expr, is_outer),
Expr::Binary { lhs, op, rhs } => {
matches!(op, B::Add | B::Sub | B::Mul | B::Div | B::IntDiv | B::Mod)
&& shape(lhs, is_outer)
&& shape(rhs, is_outer)
}
_ => false,
}
}
fn mentions_column(e: &Expr) -> bool {
match e {
Expr::Column(_) => true,
Expr::Unary { expr, .. } | Expr::Cast { expr, .. } => mentions_column(expr),
Expr::Binary { lhs, rhs, .. } => mentions_column(lhs) || mentions_column(rhs),
_ => false,
}
}
shape(e, &is_outer) && mentions_column(e)
};
let mut inner_keys: Vec<spg_sql::ast::ColumnName> = Vec::new();
let mut outer_cols: Vec<Expr> = Vec::new();
let mut rest: Vec<&Expr> = Vec::new();
for c in conjuncts {
if let Expr::Binary {
lhs,
op: BinOp::Eq,
rhs,
} = c
{
let pair = match (lhs.as_ref(), rhs.as_ref()) {
(Expr::Column(a), other) if is_inner(a) && outer_only_key(other) => {
Some((a.clone(), other.clone()))
}
(other, Expr::Column(b)) if is_inner(b) && outer_only_key(other) => {
Some((b.clone(), other.clone()))
}
_ => None,
};
if let Some((ic, oc)) = pair {
inner_keys.push(ic);
outer_cols.push(oc);
continue;
}
}
if !all_inner(c) {
return Ok(None);
}
rest.push(c);
}
if inner_keys.is_empty() {
return Ok(None); }
let mut batch = inner.clone();
batch.limit = None;
batch.offset = None;
batch.order_by = Vec::new();
batch.distinct = false;
batch.where_ = rest
.iter()
.map(|e| (*e).clone())
.reduce(|a, b| Expr::Binary {
lhs: alloc::boxed::Box::new(a),
op: BinOp::And,
rhs: alloc::boxed::Box::new(b),
});
batch.items = inner_keys
.iter()
.map(|c| SI::Expr {
expr: Expr::Column(c.clone()),
alias: None,
})
.collect();
let r = self.exec_select_cancel(&batch, cancel)?;
let QueryResult::Rows { rows, .. } = r else {
return Ok(None);
};
let n = inner_keys.len();
let mut set: alloc::collections::BTreeSet<String> = alloc::collections::BTreeSet::new();
for row in rows {
let keys = row.values.get(..n).unwrap_or(&row.values);
if keys.iter().any(|v| matches!(v, Value::Null)) {
continue;
}
set.insert(aggregate::encode_canonical_key(keys));
}
Ok(Some((outer_cols, set)))
}
}
impl Engine {
pub(crate) fn pull_up_correlated_limit_one_subqueries(
&self,
stmt: &mut SelectStatement,
) -> bool {
if stmt.from.is_none() || stmt.items.iter().any(|i| matches!(i, SelectItem::Wildcard)) {
return false;
}
let outer_aliases: alloc::collections::BTreeSet<String> = {
let from = stmt.from.as_ref().expect("from present");
let mut s = alloc::collections::BTreeSet::new();
let push = |s: &mut alloc::collections::BTreeSet<String>, t: &TableRef| {
s.insert(
t.alias
.clone()
.unwrap_or_else(|| t.name.clone())
.to_ascii_lowercase(),
);
};
push(&mut s, &from.primary);
for j in &from.joins {
push(&mut s, &j.table);
}
s
};
let outer_has_group_by = stmt.group_by.is_some() || stmt.group_by_all;
let mut new_ctes: Vec<Cte> = Vec::new();
let mut new_joins: Vec<FromJoin> = Vec::new();
let cte_seed = stmt.ctes.len();
for item in &mut stmt.items {
if let SelectItem::Expr { expr, .. } = item {
self.pull_up_walk_limit_one(
expr,
false,
&outer_aliases,
outer_has_group_by,
cte_seed,
&mut new_ctes,
&mut new_joins,
);
}
}
if new_ctes.is_empty() {
return false;
}
PULLUP_LIMIT1_FIRE_COUNT
.fetch_add(new_ctes.len() as u64, core::sync::atomic::Ordering::Relaxed);
stmt.ctes.extend(new_ctes);
stmt.from
.as_mut()
.expect("from present")
.joins
.extend(new_joins);
true
}
#[allow(clippy::too_many_arguments)]
fn pull_up_walk_limit_one(
&self,
e: &mut Expr,
in_agg: bool,
outer_aliases: &alloc::collections::BTreeSet<String>,
outer_has_group_by: bool,
cte_seed: usize,
ctes_out: &mut Vec<Cte>,
joins_out: &mut Vec<FromJoin>,
) {
match e {
Expr::ScalarSubquery(inner) => {
if let Some((cte, join, cte_col)) =
self.try_pull_up_limit_one(inner, outer_aliases, cte_seed + ctes_out.len())
{
ctes_out.push(cte);
joins_out.push(join);
let col_expr = Expr::Column(cte_col);
*e = if outer_has_group_by && !in_agg {
Expr::FunctionCall {
name: "max".into(),
args: alloc::vec![col_expr],
}
} else {
col_expr
};
}
}
Expr::FunctionCall { name, args } => {
let child = in_agg || aggregate::is_aggregate_name(name);
for a in args.iter_mut() {
self.pull_up_walk_limit_one(
a,
child,
outer_aliases,
outer_has_group_by,
cte_seed,
ctes_out,
joins_out,
);
}
}
Expr::AggregateOrdered {
call,
order_by,
filter,
..
} => {
self.pull_up_walk_limit_one(
call,
true,
outer_aliases,
outer_has_group_by,
cte_seed,
ctes_out,
joins_out,
);
for o in order_by.iter_mut() {
self.pull_up_walk_limit_one(
&mut o.expr,
true,
outer_aliases,
outer_has_group_by,
cte_seed,
ctes_out,
joins_out,
);
}
if let Some(f) = filter {
self.pull_up_walk_limit_one(
f,
true,
outer_aliases,
outer_has_group_by,
cte_seed,
ctes_out,
joins_out,
);
}
}
Expr::Binary { lhs, rhs, .. } => {
self.pull_up_walk_limit_one(
lhs,
in_agg,
outer_aliases,
outer_has_group_by,
cte_seed,
ctes_out,
joins_out,
);
self.pull_up_walk_limit_one(
rhs,
in_agg,
outer_aliases,
outer_has_group_by,
cte_seed,
ctes_out,
joins_out,
);
}
Expr::Unary { expr, .. }
| Expr::Cast { expr, .. }
| Expr::IsNull { expr, .. }
| Expr::BoolTest { expr, .. }
| Expr::FieldAccess { base: expr, .. } => {
self.pull_up_walk_limit_one(
expr,
in_agg,
outer_aliases,
outer_has_group_by,
cte_seed,
ctes_out,
joins_out,
);
}
Expr::Like { expr, pattern, .. } => {
self.pull_up_walk_limit_one(
expr,
in_agg,
outer_aliases,
outer_has_group_by,
cte_seed,
ctes_out,
joins_out,
);
self.pull_up_walk_limit_one(
pattern,
in_agg,
outer_aliases,
outer_has_group_by,
cte_seed,
ctes_out,
joins_out,
);
}
Expr::InList { expr, list, .. } => {
self.pull_up_walk_limit_one(
expr,
in_agg,
outer_aliases,
outer_has_group_by,
cte_seed,
ctes_out,
joins_out,
);
for it in list.iter_mut() {
self.pull_up_walk_limit_one(
it,
in_agg,
outer_aliases,
outer_has_group_by,
cte_seed,
ctes_out,
joins_out,
);
}
}
Expr::Case {
operand,
branches,
else_branch,
} => {
if let Some(o) = operand {
self.pull_up_walk_limit_one(
o,
in_agg,
outer_aliases,
outer_has_group_by,
cte_seed,
ctes_out,
joins_out,
);
}
for (w, t) in branches.iter_mut() {
self.pull_up_walk_limit_one(
w,
in_agg,
outer_aliases,
outer_has_group_by,
cte_seed,
ctes_out,
joins_out,
);
self.pull_up_walk_limit_one(
t,
in_agg,
outer_aliases,
outer_has_group_by,
cte_seed,
ctes_out,
joins_out,
);
}
if let Some(eb) = else_branch {
self.pull_up_walk_limit_one(
eb,
in_agg,
outer_aliases,
outer_has_group_by,
cte_seed,
ctes_out,
joins_out,
);
}
}
_ => {}
}
}
fn try_pull_up_limit_one(
&self,
inner: &SelectStatement,
outer_aliases: &alloc::collections::BTreeSet<String>,
alias_n: usize,
) -> Option<(Cte, FromJoin, ColumnName)> {
return None;
#[allow(unreachable_code)]
if !inner.ctes.is_empty()
|| !inner.unions.is_empty()
|| inner.group_by.is_some()
|| inner.group_by_all
|| inner.having.is_some()
|| inner.distinct
|| inner.offset.is_some()
|| inner.items.len() != 1
|| inner.order_by.is_empty()
{
return None;
}
match inner.limit {
Some(LimitExpr::Literal(1)) => {}
_ => return None,
}
let from = inner.from.as_ref()?;
if !from.joins.is_empty()
|| 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;
}
let inner_table = from.primary.name.clone();
let inner_alias = from
.primary
.alias
.clone()
.unwrap_or_else(|| inner_table.clone());
let is_inner = |c: &ColumnName| -> bool {
c.qualifier
.as_deref()
.is_some_and(|q| q.eq_ignore_ascii_case(&inner_alias))
};
let is_outer = |c: &ColumnName| -> bool {
c.qualifier
.as_deref()
.is_some_and(|q| outer_aliases.contains(&q.to_ascii_lowercase()))
};
let SelectItem::Expr {
expr: proj_expr,
alias: _,
} = &inner.items[0]
else {
return None;
};
if proj_has_disqualifying_shape(proj_expr, &inner_alias, outer_aliases) {
return None;
}
let where_ = inner.where_.as_ref()?;
let mut corr: Option<(String, ColumnName)> = None;
let mut non_corr: Vec<Expr> = Vec::new();
for c in reorder::split_and_conjunctions(where_) {
if let Expr::Binary {
lhs,
op: BinOp::Eq,
rhs,
} = c
&& let (Expr::Column(a), Expr::Column(b)) = (lhs.as_ref(), rhs.as_ref())
{
let pair = if is_inner(a) && is_outer(b) {
Some((a.name.clone(), b.clone()))
} else if is_inner(b) && is_outer(a) {
Some((b.name.clone(), a.clone()))
} else {
None
};
if let Some(p) = pair {
if corr.is_some() {
return None; }
corr = Some(p);
continue;
}
}
if !expr_is_all_inner(c, &inner_alias) {
return None;
}
non_corr.push(c.clone());
}
let (inner_key, outer_col) = corr?;
for ob in &inner.order_by {
if !expr_is_all_inner(&ob.expr, &inner_alias) {
return None;
}
}
if !expr_is_all_inner(proj_expr, &inner_alias) {
return None;
}
let cte_name = alloc::format!("__cl1_{alias_n}");
let jk_expr = Expr::Column(ColumnName {
qualifier: Some(inner_alias.clone()),
name: inner_key.clone(),
});
let argmax = Expr::ArraySubscript {
target: alloc::boxed::Box::new(Expr::AggregateOrdered {
call: alloc::boxed::Box::new(Expr::FunctionCall {
name: "array_agg".into(),
args: alloc::vec![proj_expr.clone()],
}),
order_by: inner.order_by.clone(),
distinct: false,
filter: None,
}),
index: alloc::boxed::Box::new(Expr::Literal(Literal::Integer(1))),
};
let body_where = if non_corr.is_empty() {
None
} else {
let mut iter = non_corr.into_iter();
let head = iter.next().expect("non_corr nonempty in this branch");
Some(iter.fold(head, |acc, p| Expr::Binary {
lhs: alloc::boxed::Box::new(acc),
op: BinOp::And,
rhs: alloc::boxed::Box::new(p),
}))
};
let body = SelectStatement {
locking: None,
ctes: Vec::new(),
distinct: false,
distinct_on: Vec::new(),
items: alloc::vec![
SelectItem::Expr {
expr: jk_expr.clone(),
alias: Some("jk".into()),
},
SelectItem::Expr {
expr: argmax,
alias: Some("pj".into()),
},
],
from: Some(from.clone()),
where_: body_where,
group_by: Some(alloc::vec![jk_expr]),
group_by_all: false,
having: None,
unions: Vec::new(),
order_by: Vec::new(),
limit: None,
offset: None,
limit_with_ties: false,
window_check_exprs: Vec::new(),
};
let cte = Cte {
name: cte_name.clone(),
body: spg_sql::ast::CteBody::Select(body),
recursive: false,
column_overrides: Vec::new(),
search: None,
cycle: None,
};
let join = FromJoin {
kind: JoinKind::Left,
table: TableRef {
name: cte_name.clone(),
alias: None,
only: false,
as_of_segment: None,
unnest_expr: None,
unnest_column_aliases: Vec::new(),
with_ordinality: false,
generate_series_args: None,
lateral_subquery: None,
jsonb_each_text_arg: None,
table_fn_call: None,
rows_from: None,
json_table: None,
scalar_fn_item: false,
},
on: Some(Expr::Binary {
lhs: alloc::boxed::Box::new(Expr::Column(ColumnName {
qualifier: Some(cte_name.clone()),
name: "jk".into(),
})),
op: BinOp::Eq,
rhs: alloc::boxed::Box::new(Expr::Column(outer_col)),
}),
using_cols: None,
natural: false,
};
let repl = ColumnName {
qualifier: Some(cte_name),
name: "pj".into(),
};
Some((cte, join, repl))
}
pub(crate) fn pull_up_unique_correlated_agg_subqueries(
&self,
stmt: &mut SelectStatement,
) -> bool {
if stmt.from.is_none() || stmt.items.iter().any(|i| matches!(i, SelectItem::Wildcard)) {
return false;
}
let outer_aliases: alloc::collections::BTreeSet<String> = {
let from = stmt.from.as_ref().expect("from present");
let mut s = alloc::collections::BTreeSet::new();
let push = |s: &mut alloc::collections::BTreeSet<String>, t: &TableRef| {
s.insert(
t.alias
.clone()
.unwrap_or_else(|| t.name.clone())
.to_ascii_lowercase(),
);
};
push(&mut s, &from.primary);
for j in &from.joins {
push(&mut s, &j.table);
}
s
};
let mut new_joins: Vec<FromJoin> = Vec::new();
for item in &mut stmt.items {
if let SelectItem::Expr { expr, .. } = item {
self.pull_up_walk(expr, false, &outer_aliases, &mut new_joins);
}
}
if new_joins.is_empty() {
return false;
}
stmt.from
.as_mut()
.expect("from present")
.joins
.extend(new_joins);
true
}
fn pull_up_walk(
&self,
e: &mut Expr,
in_agg: bool,
outer_aliases: &alloc::collections::BTreeSet<String>,
joins_out: &mut Vec<FromJoin>,
) {
match e {
Expr::ScalarSubquery(inner) => {
if in_agg
&& let Some((join, col)) =
self.try_pull_up_join(inner, outer_aliases, joins_out.len())
{
joins_out.push(join);
*e = Expr::Column(col);
}
}
Expr::FunctionCall { name, args } => {
let child = in_agg || aggregate::is_aggregate_name(name);
for a in args.iter_mut() {
self.pull_up_walk(a, child, outer_aliases, joins_out);
}
}
Expr::AggregateOrdered {
call,
order_by,
filter,
..
} => {
self.pull_up_walk(call, true, outer_aliases, joins_out);
for o in order_by.iter_mut() {
self.pull_up_walk(&mut o.expr, true, outer_aliases, joins_out);
}
if let Some(f) = filter {
self.pull_up_walk(f, true, outer_aliases, joins_out);
}
}
Expr::Binary { lhs, rhs, .. } => {
self.pull_up_walk(lhs, in_agg, outer_aliases, joins_out);
self.pull_up_walk(rhs, in_agg, outer_aliases, joins_out);
}
Expr::Unary { expr, .. }
| Expr::Cast { expr, .. }
| Expr::IsNull { expr, .. }
| Expr::BoolTest { expr, .. }
| Expr::FieldAccess { base: expr, .. } => {
self.pull_up_walk(expr, in_agg, outer_aliases, joins_out);
}
Expr::Like { expr, pattern, .. } => {
self.pull_up_walk(expr, in_agg, outer_aliases, joins_out);
self.pull_up_walk(pattern, in_agg, outer_aliases, joins_out);
}
Expr::InList { expr, list, .. } => {
self.pull_up_walk(expr, in_agg, outer_aliases, joins_out);
for it in list.iter_mut() {
self.pull_up_walk(it, in_agg, outer_aliases, joins_out);
}
}
Expr::Case {
operand,
branches,
else_branch,
} => {
if let Some(o) = operand {
self.pull_up_walk(o, in_agg, outer_aliases, joins_out);
}
for (w, t) in branches.iter_mut() {
self.pull_up_walk(w, in_agg, outer_aliases, joins_out);
self.pull_up_walk(t, in_agg, outer_aliases, joins_out);
}
if let Some(eb) = else_branch {
self.pull_up_walk(eb, in_agg, outer_aliases, joins_out);
}
}
_ => {}
}
}
fn try_pull_up_join(
&self,
inner: &SelectStatement,
outer_aliases: &alloc::collections::BTreeSet<String>,
alias_n: usize,
) -> Option<(FromJoin, ColumnName)> {
if !inner.ctes.is_empty()
|| !inner.unions.is_empty()
|| inner.group_by.is_some()
|| inner.having.is_some()
|| inner.distinct
|| !inner.order_by.is_empty()
|| inner.limit.is_some()
|| inner.offset.is_some()
|| inner.items.len() != 1
{
return None;
}
let from = inner.from.as_ref()?;
if !from.joins.is_empty()
|| 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;
}
let inner_table = from.primary.name.clone();
let inner_alias = from
.primary
.alias
.clone()
.unwrap_or_else(|| inner_table.clone());
let is_inner = |c: &ColumnName| -> bool {
c.qualifier
.as_deref()
.is_some_and(|q| q.eq_ignore_ascii_case(&inner_alias))
};
let is_outer = |c: &ColumnName| -> bool {
c.qualifier
.as_deref()
.is_some_and(|q| outer_aliases.contains(&q.to_ascii_lowercase()))
};
let SelectItem::Expr { expr: out_expr, .. } = &inner.items[0] else {
return None;
};
let Expr::Column(out_col) = out_expr else {
return None;
};
if !is_inner(out_col) {
return None;
}
let w = inner.where_.as_ref()?;
let mut corr: Option<(String, ColumnName)> = None;
let mut rest: Vec<Expr> = Vec::new();
for c in reorder::split_and_conjunctions(w) {
if let Expr::Binary {
lhs,
op: BinOp::Eq,
rhs,
} = c
&& let (Expr::Column(a), Expr::Column(b)) = (lhs.as_ref(), rhs.as_ref())
{
let pair = if is_inner(a) && is_outer(b) {
Some((a.name.clone(), b.clone()))
} else if is_inner(b) && is_outer(a) {
Some((b.name.clone(), a.clone()))
} else {
None
};
if let Some(p) = pair {
if corr.is_some() {
return None; }
corr = Some(p);
continue;
}
}
if !expr_is_all_inner(c, &inner_alias) {
return None;
}
rest.push(c.clone());
}
let (inner_key, outer_col) = corr?;
if !self.column_is_single_unique(&inner_table, &inner_key) {
return None;
}
let fresh = alloc::format!("__plj_{alias_n}");
let key_eq = Expr::Binary {
lhs: alloc::boxed::Box::new(Expr::Column(ColumnName {
qualifier: Some(fresh.clone()),
name: inner_key,
})),
op: BinOp::Eq,
rhs: alloc::boxed::Box::new(Expr::Column(outer_col)),
};
let on = rest
.into_iter()
.map(|mut e| {
rename_qualifier(&mut e, &inner_alias, &fresh);
e
})
.fold(key_eq, |acc, pred| Expr::Binary {
lhs: alloc::boxed::Box::new(acc),
op: BinOp::And,
rhs: alloc::boxed::Box::new(pred),
});
let join = FromJoin {
kind: JoinKind::Left,
table: TableRef {
name: inner_table,
alias: Some(fresh.clone()),
only: false,
as_of_segment: None,
unnest_expr: None,
unnest_column_aliases: Vec::new(),
with_ordinality: false,
generate_series_args: None,
lateral_subquery: None,
jsonb_each_text_arg: None,
table_fn_call: None,
rows_from: None,
json_table: None,
scalar_fn_item: false,
},
on: Some(on),
using_cols: None,
natural: false,
};
let repl = ColumnName {
qualifier: Some(fresh),
name: out_col.name.clone(),
};
Some((join, repl))
}
pub(crate) fn pull_up_exists_sublinks(&self, stmt: &mut SelectStatement) -> bool {
if stmt.from.is_none() {
return false;
}
let Some(where_expr) = stmt.where_.take() else {
return false;
};
let mut collision_names: alloc::collections::BTreeSet<String> =
alloc::collections::BTreeSet::new();
for c in reorder::split_and_conjunctions(&where_expr) {
let inner_subq: Option<&SelectStatement> = match c {
Expr::Exists { subquery, .. } => Some(subquery.as_ref()),
Expr::Unary {
op: UnOp::Not,
expr,
} => match expr.as_ref() {
Expr::Exists { subquery, .. } => Some(subquery.as_ref()),
_ => None,
},
_ => None,
};
let Some(inner) = inner_subq else { continue };
let Some(from) = &inner.from else { continue };
if !from.joins.is_empty() {
continue;
}
let Some(t) = self.active_catalog().get(&from.primary.name) else {
continue;
};
for col in &t.schema().columns {
collision_names.insert(col.name.to_ascii_lowercase());
}
}
let mut where_expr = where_expr;
if !collision_names.is_empty() {
let from = stmt.from.as_ref().expect("from present");
let outer_tables: Vec<(String, String)> = {
let mut v = Vec::new();
let collect = |v: &mut Vec<(String, String)>, t: &TableRef| {
let alias = t.alias.clone().unwrap_or_else(|| t.name.clone());
v.push((alias, t.name.clone()));
};
collect(&mut v, &from.primary);
for j in &from.joins {
collect(&mut v, &j.table);
}
v
};
let mut owner: alloc::collections::BTreeMap<String, String> =
alloc::collections::BTreeMap::new();
for col_lc in &collision_names {
let mut matches: Vec<String> = Vec::new();
for (alias, tname) in &outer_tables {
let Some(t) = self.active_catalog().get(tname) else {
continue;
};
if t.schema()
.columns
.iter()
.any(|c| c.name.eq_ignore_ascii_case(col_lc))
{
matches.push(alias.clone());
}
}
if matches.len() == 1 {
owner.insert(col_lc.clone(), matches.remove(0));
}
}
if !owner.is_empty() {
disambiguate_stmt_unqualified_columns(stmt, &owner);
disambiguate_expr_unqualified_columns(&mut where_expr, &owner);
}
}
let outer_aliases: alloc::collections::BTreeSet<String> = {
let from = stmt.from.as_ref().expect("from present");
let mut s = alloc::collections::BTreeSet::new();
let push = |s: &mut alloc::collections::BTreeSet<String>, t: &TableRef| {
s.insert(
t.alias
.clone()
.unwrap_or_else(|| t.name.clone())
.to_ascii_lowercase(),
);
};
push(&mut s, &from.primary);
for j in &from.joins {
push(&mut s, &j.table);
}
s
};
let outer_tables: alloc::collections::BTreeMap<String, String> = {
let from = stmt.from.as_ref().expect("from present");
let mut m = alloc::collections::BTreeMap::new();
let push = |m: &mut alloc::collections::BTreeMap<String, String>, t: &TableRef| {
m.insert(
t.alias
.clone()
.unwrap_or_else(|| t.name.clone())
.to_ascii_lowercase(),
t.name.clone(),
);
};
push(&mut m, &from.primary);
for j in &from.joins {
push(&mut m, &j.table);
}
m
};
let conjuncts = reorder::split_and_conjunctions(&where_expr);
let mut survivors: Vec<Expr> = Vec::new();
let mut new_joins: Vec<FromJoin> = Vec::new();
let mut rewrote_any = false;
for c in conjuncts {
let parsed: Option<(&SelectStatement, bool)> = match c {
Expr::Exists { subquery, negated } => Some((subquery.as_ref(), *negated)),
Expr::Unary {
op: UnOp::Not,
expr,
} => match expr.as_ref() {
Expr::Exists { subquery, negated } => Some((subquery.as_ref(), !*negated)),
_ => None,
},
_ => None,
};
if let Some((subquery, neg)) = parsed {
let outer_has_limit = stmt.limit.is_some();
let try_in_first = outer_has_limit;
let mut consumed = false;
if try_in_first
&& let Some(rewritten) =
self.try_pull_up_exists_as_in(subquery, neg, &outer_aliases)
{
survivors.push(rewritten);
consumed = true;
}
if !consumed
&& let Some((join, residual)) = self.try_pull_up_exists_sublink(
subquery,
neg,
&outer_aliases,
&outer_tables,
new_joins.len(),
)
{
new_joins.push(join);
if let Some(r) = residual {
survivors.push(r);
}
consumed = true;
}
if !consumed
&& !try_in_first
&& let Some(rewritten) =
self.try_pull_up_exists_as_in(subquery, neg, &outer_aliases)
{
survivors.push(rewritten);
consumed = true;
}
if consumed {
rewrote_any = true;
continue;
}
}
survivors.push(c.clone());
}
if !rewrote_any {
stmt.where_ = Some(where_expr);
return false;
}
EXISTS_PULLUP_FIRE_COUNT.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
if !new_joins.is_empty() {
stmt.from
.as_mut()
.expect("from present")
.joins
.extend(new_joins);
}
stmt.where_ = survivors.into_iter().reduce(|a, b| Expr::Binary {
lhs: alloc::boxed::Box::new(a),
op: BinOp::And,
rhs: alloc::boxed::Box::new(b),
});
true
}
fn try_pull_up_exists_as_in(
&self,
inner: &SelectStatement,
negated: bool,
outer_aliases: &alloc::collections::BTreeSet<String>,
) -> Option<Expr> {
if !inner.ctes.is_empty()
|| !inner.unions.is_empty()
|| inner.group_by.is_some()
|| inner.having.is_some()
|| inner.distinct
|| !inner.order_by.is_empty()
|| inner.limit.is_some()
|| inner.offset.is_some()
{
return None;
}
let from = inner.from.as_ref()?;
if !from.joins.is_empty()
|| 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;
}
let inner_table = from.primary.name.clone();
let inner_alias = from
.primary
.alias
.clone()
.unwrap_or_else(|| inner_table.clone());
let is_inner = |c: &ColumnName| -> bool {
c.qualifier
.as_deref()
.is_some_and(|q| q.eq_ignore_ascii_case(&inner_alias))
};
let is_outer = |c: &ColumnName| -> bool {
c.qualifier
.as_deref()
.is_some_and(|q| outer_aliases.contains(&q.to_ascii_lowercase()))
};
let w = inner.where_.as_ref()?;
let mut corr: Option<(String, ColumnName)> = None;
let mut rest: Vec<Expr> = Vec::new();
for c in reorder::split_and_conjunctions(w) {
if let Expr::Binary {
lhs,
op: BinOp::Eq,
rhs,
} = c
&& let (Expr::Column(a), Expr::Column(b)) = (lhs.as_ref(), rhs.as_ref())
{
let pair = if is_inner(a) && is_outer(b) {
Some((a.name.clone(), b.clone()))
} else if is_inner(b) && is_outer(a) {
Some((b.name.clone(), a.clone()))
} else {
None
};
if let Some(p) = pair {
if corr.is_some() {
return None;
}
corr = Some(p);
continue;
}
}
if !expr_is_all_inner(c, &inner_alias) {
return None;
}
rest.push(c.clone());
}
let (inner_key, outer_col) = corr?;
if negated && !self.column_is_not_null(&inner_table, &inner_key) {
return None;
}
let mut rewritten = inner.clone();
rewritten.limit = None;
rewritten.offset = None;
rewritten.order_by = Vec::new();
rewritten.distinct = false;
rewritten.where_ = rest.into_iter().reduce(|a, b| Expr::Binary {
lhs: alloc::boxed::Box::new(a),
op: BinOp::And,
rhs: alloc::boxed::Box::new(b),
});
rewritten.items = alloc::vec![SelectItem::Expr {
expr: Expr::Column(ColumnName {
qualifier: Some(inner_alias),
name: inner_key,
}),
alias: None,
}];
Some(Expr::InSubquery {
expr: alloc::boxed::Box::new(Expr::Column(outer_col)),
subquery: alloc::boxed::Box::new(rewritten),
negated,
})
}
fn try_pull_up_exists_sublink(
&self,
inner: &SelectStatement,
negated: bool,
outer_aliases: &alloc::collections::BTreeSet<String>,
outer_tables: &alloc::collections::BTreeMap<String, String>,
alias_n: usize,
) -> Option<(FromJoin, Option<Expr>)> {
EXISTS_PULLUP_CANDIDATE_COUNT.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
if !inner.ctes.is_empty()
|| !inner.unions.is_empty()
|| inner.group_by.is_some()
|| inner.having.is_some()
|| inner.distinct
|| !inner.order_by.is_empty()
|| inner.limit.is_some()
|| inner.offset.is_some()
{
EXISTS_PULLUP_BAIL_INNER_SHAPE.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
return None;
}
let from = inner.from.as_ref()?;
if !from.joins.is_empty()
|| 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()
{
EXISTS_PULLUP_BAIL_INNER_FROM.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
return None;
}
let inner_table = from.primary.name.clone();
let inner_alias = from
.primary
.alias
.clone()
.unwrap_or_else(|| inner_table.clone());
let is_inner = |c: &ColumnName| -> bool {
c.qualifier
.as_deref()
.is_some_and(|q| q.eq_ignore_ascii_case(&inner_alias))
};
let is_outer = |c: &ColumnName| -> bool {
c.qualifier
.as_deref()
.is_some_and(|q| outer_aliases.contains(&q.to_ascii_lowercase()))
};
let Some(w) = inner.where_.as_ref() else {
EXISTS_PULLUP_BAIL_NO_WHERE.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
return None;
};
let inner_col_is_int = |name: &str| -> bool {
self.active_catalog().get(&inner_table).is_some_and(|t| {
t.schema().columns.iter().any(|cs| {
cs.name.eq_ignore_ascii_case(name)
&& matches!(
cs.ty,
spg_storage::DataType::Int
| spg_storage::DataType::BigInt
| spg_storage::DataType::SmallInt
)
})
})
};
enum InnerHalf {
Col(String),
IntExpr(Expr),
}
fn inner_int_only_expr(
e: &Expr,
is_inner: &dyn Fn(&ColumnName) -> bool,
inner_col_is_int: &dyn Fn(&str) -> bool,
) -> bool {
match e {
Expr::Column(c) => is_inner(c) && inner_col_is_int(&c.name),
Expr::Literal(spg_sql::ast::Literal::Integer(_)) => true,
Expr::Binary { lhs, op, rhs } => {
matches!(op, BinOp::Add | BinOp::Sub | BinOp::Mul)
&& inner_int_only_expr(lhs, is_inner, inner_col_is_int)
&& inner_int_only_expr(rhs, is_inner, inner_col_is_int)
}
_ => false,
}
}
fn first_inner_col(e: &Expr) -> Option<String> {
match e {
Expr::Column(c) => Some(c.name.clone()),
Expr::Binary { lhs, rhs, .. } => {
first_inner_col(lhs).or_else(|| first_inner_col(rhs))
}
_ => None,
}
}
let mut corr_pairs: Vec<(InnerHalf, Expr)> = Vec::new();
let mut rest: Vec<Expr> = Vec::new();
for c in reorder::split_and_conjunctions(w) {
if let Expr::Binary {
lhs,
op: BinOp::Eq,
rhs,
} = c
{
let pair = match (lhs.as_ref(), rhs.as_ref()) {
(Expr::Column(a), Expr::Column(b)) if is_inner(a) && is_outer(b) => {
Some((InnerHalf::Col(a.name.clone()), Expr::Column(b.clone())))
}
(Expr::Column(a), Expr::Column(b)) if is_inner(b) && is_outer(a) => {
Some((InnerHalf::Col(b.name.clone()), Expr::Column(a.clone())))
}
(Expr::Column(a), e)
if is_inner(a)
&& !matches!(e, Expr::Column(_))
&& inner_col_is_int(&a.name)
&& outer_int_only_expr(e, outer_aliases, outer_tables, self) =>
{
Some((InnerHalf::Col(a.name.clone()), e.clone()))
}
(e, Expr::Column(a))
if is_inner(a)
&& !matches!(e, Expr::Column(_))
&& inner_col_is_int(&a.name)
&& outer_int_only_expr(e, outer_aliases, outer_tables, self) =>
{
Some((InnerHalf::Col(a.name.clone()), e.clone()))
}
(Expr::Column(o), e)
if is_outer(o)
&& !matches!(e, Expr::Column(_))
&& outer_int_only_expr(
&Expr::Column(o.clone()),
outer_aliases,
outer_tables,
self,
)
&& inner_int_only_expr(e, &is_inner, &inner_col_is_int)
&& first_inner_col(e).is_some() =>
{
Some((InnerHalf::IntExpr(e.clone()), Expr::Column(o.clone())))
}
(e, Expr::Column(o))
if is_outer(o)
&& !matches!(e, Expr::Column(_))
&& outer_int_only_expr(
&Expr::Column(o.clone()),
outer_aliases,
outer_tables,
self,
)
&& inner_int_only_expr(e, &is_inner, &inner_col_is_int)
&& first_inner_col(e).is_some() =>
{
Some((InnerHalf::IntExpr(e.clone()), Expr::Column(o.clone())))
}
_ => None,
};
if let Some(p) = pair {
corr_pairs.push(p);
continue;
}
}
if !expr_is_all_inner(c, &inner_alias) {
EXISTS_PULLUP_BAIL_RESIDUAL_NOT_INNER
.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
return None;
}
rest.push(c.clone());
}
if corr_pairs.is_empty() {
EXISTS_PULLUP_BAIL_NO_CORR.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
return None;
}
if corr_pairs.len() > 1
&& EXISTS_PULLUP_MULTICOL_DISABLE.load(core::sync::atomic::Ordering::Relaxed)
{
EXISTS_PULLUP_BAIL_MULTICOL_DISABLED
.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
return None;
}
let fresh = alloc::format!("__exsj_{alias_n}");
let inner_half_expr = |ih: &InnerHalf| -> Expr {
match ih {
InnerHalf::Col(name) => Expr::Column(ColumnName {
qualifier: Some(fresh.clone()),
name: name.clone(),
}),
InnerHalf::IntExpr(e) => {
let mut e = e.clone();
rename_qualifier(&mut e, &inner_alias, &fresh);
e
}
}
};
let mut on_iter = corr_pairs.iter().map(|(ik, oc)| Expr::Binary {
lhs: alloc::boxed::Box::new(inner_half_expr(ik)),
op: BinOp::Eq,
rhs: alloc::boxed::Box::new(oc.clone()),
});
let first_key_eq = on_iter
.next()
.expect("corr_pairs non-empty post `is_empty()` gate");
let on = rest
.into_iter()
.map(|mut e| {
rename_qualifier(&mut e, &inner_alias, &fresh);
e
})
.chain(on_iter)
.fold(first_key_eq, |acc, pred| Expr::Binary {
lhs: alloc::boxed::Box::new(acc),
op: BinOp::And,
rhs: alloc::boxed::Box::new(pred),
});
let join = FromJoin {
kind: if negated {
JoinKind::Left
} else {
JoinKind::Semi
},
table: TableRef {
name: inner_table,
alias: Some(fresh.clone()),
only: false,
as_of_segment: None,
unnest_expr: None,
unnest_column_aliases: Vec::new(),
with_ordinality: false,
generate_series_args: None,
lateral_subquery: None,
jsonb_each_text_arg: None,
table_fn_call: None,
rows_from: None,
json_table: None,
scalar_fn_item: false,
},
on: Some(on),
using_cols: None,
natural: false,
};
let residual = if negated {
let probe_key = match &corr_pairs[0].0 {
InnerHalf::Col(name) => name.clone(),
InnerHalf::IntExpr(e) => {
first_inner_col(e).expect("IntExpr admitted only with an inner column")
}
};
Some(Expr::IsNull {
expr: alloc::boxed::Box::new(Expr::Column(ColumnName {
qualifier: Some(fresh),
name: probe_key,
})),
negated: false,
})
} else {
None
};
Some((join, residual))
}
fn column_is_not_null(&self, table: &str, col: &str) -> bool {
let Some(t) = self.active_catalog().get(table) else {
return false;
};
let sch = t.schema();
if sch
.columns
.iter()
.find(|c| c.name.eq_ignore_ascii_case(col))
.is_some_and(|c| !c.nullable)
{
return true;
}
let Some(pos) = sch
.columns
.iter()
.position(|c| c.name.eq_ignore_ascii_case(col))
else {
return false;
};
sch.uniqueness_constraints
.iter()
.any(|u| u.is_primary_key && u.columns.as_slice() == [pos])
}
fn column_is_single_unique(&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 proj_has_disqualifying_shape(
e: &Expr,
inner_alias: &str,
outer_aliases: &alloc::collections::BTreeSet<String>,
) -> bool {
match e {
Expr::AggregateOrdered { .. }
| Expr::WindowFunction { .. }
| Expr::ScalarSubquery(_)
| Expr::Exists { .. } => true,
Expr::FunctionCall { name, args } => {
if aggregate::is_aggregate_name(name) {
return true;
}
args.iter()
.any(|a| proj_has_disqualifying_shape(a, inner_alias, outer_aliases))
}
Expr::Column(c) => {
if let Some(q) = c.qualifier.as_deref() {
outer_aliases.contains(&q.to_ascii_lowercase())
&& !q.eq_ignore_ascii_case(inner_alias)
} else {
false
}
}
Expr::Binary { lhs, rhs, .. } => {
proj_has_disqualifying_shape(lhs, inner_alias, outer_aliases)
|| proj_has_disqualifying_shape(rhs, inner_alias, outer_aliases)
}
Expr::Unary { expr, .. }
| Expr::Cast { expr, .. }
| Expr::IsNull { expr, .. }
| Expr::BoolTest { expr, .. }
| Expr::FieldAccess { base: expr, .. } => {
proj_has_disqualifying_shape(expr, inner_alias, outer_aliases)
}
Expr::Like { expr, pattern, .. } => {
proj_has_disqualifying_shape(expr, inner_alias, outer_aliases)
|| proj_has_disqualifying_shape(pattern, inner_alias, outer_aliases)
}
Expr::InList { expr, list, .. } => {
proj_has_disqualifying_shape(expr, inner_alias, outer_aliases)
|| list
.iter()
.any(|it| proj_has_disqualifying_shape(it, inner_alias, outer_aliases))
}
Expr::Case {
operand,
branches,
else_branch,
} => {
operand
.as_ref()
.is_some_and(|o| proj_has_disqualifying_shape(o, inner_alias, outer_aliases))
|| branches.iter().any(|(w, t)| {
proj_has_disqualifying_shape(w, inner_alias, outer_aliases)
|| proj_has_disqualifying_shape(t, inner_alias, outer_aliases)
})
|| else_branch
.as_ref()
.is_some_and(|b| proj_has_disqualifying_shape(b, inner_alias, outer_aliases))
}
Expr::ArraySubscript { target, index } => {
proj_has_disqualifying_shape(target, inner_alias, outer_aliases)
|| proj_has_disqualifying_shape(index, inner_alias, outer_aliases)
}
_ => false,
}
}
fn disambiguate_stmt_unqualified_columns(
stmt: &mut SelectStatement,
owner: &alloc::collections::BTreeMap<String, String>,
) {
for item in &mut stmt.items {
if let SelectItem::Expr { expr, .. } = item {
disambiguate_expr_unqualified_columns(expr, owner);
}
}
if let Some(from) = &mut stmt.from {
for j in &mut from.joins {
if let Some(on) = &mut j.on {
disambiguate_expr_unqualified_columns(on, owner);
}
}
}
if let Some(g) = &mut stmt.group_by {
for e in g.iter_mut() {
disambiguate_expr_unqualified_columns(e, owner);
}
}
if let Some(h) = &mut stmt.having {
disambiguate_expr_unqualified_columns(h, owner);
}
for ob in &mut stmt.order_by {
disambiguate_expr_unqualified_columns(&mut ob.expr, owner);
}
}
fn disambiguate_expr_unqualified_columns(
e: &mut Expr,
owner: &alloc::collections::BTreeMap<String, String>,
) {
match e {
Expr::Column(c) => {
if c.qualifier.is_none()
&& let Some(alias) = owner.get(&c.name.to_ascii_lowercase())
{
c.qualifier = Some(alias.clone());
}
}
Expr::Binary { lhs, rhs, .. } => {
disambiguate_expr_unqualified_columns(lhs, owner);
disambiguate_expr_unqualified_columns(rhs, owner);
}
Expr::Unary { expr, .. }
| Expr::Cast { expr, .. }
| Expr::IsNull { expr, .. }
| Expr::BoolTest { expr, .. }
| Expr::FieldAccess { base: expr, .. } => {
disambiguate_expr_unqualified_columns(expr, owner);
}
Expr::FunctionCall { args, .. } => {
for a in args.iter_mut() {
disambiguate_expr_unqualified_columns(a, owner);
}
}
Expr::AggregateOrdered {
call,
order_by,
filter,
..
} => {
disambiguate_expr_unqualified_columns(call, owner);
for ob in order_by.iter_mut() {
disambiguate_expr_unqualified_columns(&mut ob.expr, owner);
}
if let Some(f) = filter {
disambiguate_expr_unqualified_columns(f, owner);
}
}
Expr::Like { expr, pattern, .. } => {
disambiguate_expr_unqualified_columns(expr, owner);
disambiguate_expr_unqualified_columns(pattern, owner);
}
Expr::InList { expr, list, .. } => {
disambiguate_expr_unqualified_columns(expr, owner);
for it in list.iter_mut() {
disambiguate_expr_unqualified_columns(it, owner);
}
}
Expr::Case {
operand,
branches,
else_branch,
} => {
if let Some(o) = operand {
disambiguate_expr_unqualified_columns(o, owner);
}
for (w, t) in branches.iter_mut() {
disambiguate_expr_unqualified_columns(w, owner);
disambiguate_expr_unqualified_columns(t, owner);
}
if let Some(eb) = else_branch {
disambiguate_expr_unqualified_columns(eb, owner);
}
}
Expr::ArraySubscript { target, index } => {
disambiguate_expr_unqualified_columns(target, owner);
disambiguate_expr_unqualified_columns(index, owner);
}
_ => {}
}
}
fn outer_int_only_expr(
e: &Expr,
outer_aliases: &alloc::collections::BTreeSet<String>,
outer_tables: &alloc::collections::BTreeMap<String, String>,
engine: &Engine,
) -> bool {
match e {
Expr::Column(c) => {
let Some(q) = c.qualifier.as_deref() else {
return false;
};
let q = q.to_ascii_lowercase();
if !outer_aliases.contains(&q) {
return false;
}
let Some(tname) = outer_tables.get(&q) else {
return false;
};
engine.active_catalog().get(tname).is_some_and(|t| {
t.schema().columns.iter().any(|cs| {
cs.name.eq_ignore_ascii_case(&c.name)
&& matches!(
cs.ty,
spg_storage::DataType::Int
| spg_storage::DataType::BigInt
| spg_storage::DataType::SmallInt
)
})
})
}
Expr::Literal(spg_sql::ast::Literal::Integer(_)) => true,
Expr::Binary { lhs, op, rhs } => {
matches!(op, BinOp::Add | BinOp::Sub | BinOp::Mul)
&& outer_int_only_expr(lhs, outer_aliases, outer_tables, engine)
&& outer_int_only_expr(rhs, outer_aliases, outer_tables, engine)
}
_ => false,
}
}
fn expr_is_all_inner(e: &Expr, inner_alias: &str) -> bool {
let mut cols: Vec<ColumnName> = Vec::new();
let mut subs: Vec<&SelectStatement> = Vec::new();
visit_expr_columns_and_subqueries(e, &mut |c| cols.push(c.clone()), &mut |s| subs.push(s));
subs.is_empty()
&& cols.iter().all(|c| {
c.qualifier
.as_deref()
.is_some_and(|q| q.eq_ignore_ascii_case(inner_alias))
})
}
fn rename_qualifier(e: &mut Expr, from: &str, to: &str) {
match e {
Expr::Column(c) => {
if c.qualifier
.as_deref()
.is_some_and(|q| q.eq_ignore_ascii_case(from))
{
c.qualifier = Some(to.into());
}
}
Expr::Binary { lhs, rhs, .. } => {
rename_qualifier(lhs, from, to);
rename_qualifier(rhs, from, to);
}
Expr::Unary { expr, .. }
| Expr::Cast { expr, .. }
| Expr::IsNull { expr, .. }
| Expr::BoolTest { expr, .. }
| Expr::FieldAccess { base: expr, .. } => {
rename_qualifier(expr, from, to);
}
Expr::FunctionCall { args, .. } => {
for a in args.iter_mut() {
rename_qualifier(a, from, to);
}
}
Expr::Like { expr, pattern, .. } => {
rename_qualifier(expr, from, to);
rename_qualifier(pattern, from, to);
}
Expr::InList { expr, list, .. } => {
rename_qualifier(expr, from, to);
for it in list.iter_mut() {
rename_qualifier(it, from, to);
}
}
Expr::Case {
operand,
branches,
else_branch,
} => {
if let Some(o) = operand {
rename_qualifier(o, from, to);
}
for (w, t) in branches.iter_mut() {
rename_qualifier(w, from, to);
rename_qualifier(t, from, to);
}
if let Some(eb) = else_branch {
rename_qualifier(eb, from, to);
}
}
_ => {}
}
}
fn is_correlation_error(e: &EngineError) -> bool {
matches!(
e,
EngineError::Eval(
eval::EvalError::ColumnNotFound { .. } | eval::EvalError::UnknownQualifier { .. }
)
)
}
#[derive(Debug, Clone)]
pub struct ScalarPkProbeFastPath {
pub outer_pos: usize,
pub inner_table_name: String,
pub inner_pos: usize,
pub table_idx: usize,
}
impl ScalarPkProbeFastPath {
pub fn probe(&self, row: &Row<'static>) -> Value<'static> {
let outer_int = match row.values.get(self.outer_pos) {
Some(Value::BigInt(n)) => *n,
Some(Value::Int(n)) => i64::from(*n),
Some(Value::SmallInt(n)) => i64::from(*n),
Some(Value::Null) | None => return Value::BigInt(0),
_ => return Value::BigInt(0),
};
SCALARSQ_PK_PROBE_PLAN_OUTER_INT.store(outer_int, core::sync::atomic::Ordering::Relaxed);
SCALARSQ_PK_PROBE_PLAN_FIRED.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
Value::BigInt(0)
}
}
pub static SCALARSQ_PK_PROBE_PLAN_FIRED: core::sync::atomic::AtomicU64 =
core::sync::atomic::AtomicU64::new(0);
pub static SCALARSQ_PK_PROBE_PLAN_OUTER_INT: core::sync::atomic::AtomicI64 =
core::sync::atomic::AtomicI64::new(0);
impl Engine {
pub(crate) fn probe_with_pk_fast_path(
&self,
plan: &ScalarPkProbeFastPath,
row: &Row<'static>,
) -> Value<'static> {
let outer_int = match row.values.get(plan.outer_pos) {
Some(Value::BigInt(n)) => *n,
Some(Value::Int(n)) => i64::from(*n),
Some(Value::SmallInt(n)) => i64::from(*n),
Some(Value::Null) | None => return Value::BigInt(0),
_ => return Value::BigInt(0),
};
let Some(inner_table) = self.active_catalog().tables_at(plan.table_idx) else {
return Value::BigInt(0);
};
let Some(idx) = inner_table.index_on(plan.inner_pos) else {
return Value::BigInt(0);
};
let Some(key) = spg_storage::IndexKey::from_value(&Value::BigInt(outer_int)) else {
return Value::BigInt(0);
};
let hit = !idx.lookup_eq(&key).is_empty();
SCALARSQ_PK_PROBE_FIRED.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
Value::BigInt(i64::from(hit))
}
pub(crate) fn analyse_scalar_count_pk_eq_probe(
&self,
inner: &SelectStatement,
outer_schema: &[spg_storage::ColumnSchema],
outer_alias: &str,
) -> Option<ScalarPkProbeFastPath> {
use spg_sql::ast::{BinOp, ColumnName, SelectItem};
if !inner.ctes.is_empty()
|| !inner.unions.is_empty()
|| inner.group_by.is_some()
|| inner.having.is_some()
|| inner.distinct
|| !inner.order_by.is_empty()
|| inner.limit.is_some()
|| inner.offset.is_some()
|| inner.items.len() != 1
{
return None;
}
let SelectItem::Expr { expr, .. } = &inner.items[0] else {
return None;
};
let is_count_shape = match expr {
Expr::FunctionCall { name, args } => {
(name.eq_ignore_ascii_case("count_star") && args.is_empty())
|| name.eq_ignore_ascii_case("count")
}
_ => false,
};
if !is_count_shape {
return None;
}
let from = inner.from.as_ref()?;
if !from.joins.is_empty()
|| 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;
}
let inner_table_name = from.primary.name.clone();
let inner_alias = from
.primary
.alias
.as_deref()
.unwrap_or(inner_table_name.as_str());
let where_expr = inner.where_.as_ref()?;
let Expr::Binary {
lhs,
op: BinOp::Eq,
rhs,
} = where_expr
else {
return None;
};
let (Expr::Column(a), Expr::Column(b)) = (lhs.as_ref(), rhs.as_ref()) else {
return None;
};
let pick = |x: &ColumnName, y: &ColumnName| -> Option<(String, ColumnName)> {
if x.qualifier
.as_deref()
.is_some_and(|q| q.eq_ignore_ascii_case(inner_alias))
{
Some((x.name.clone(), y.clone()))
} else {
None
}
};
let (inner_col_name, outer_col) = pick(a, b).or_else(|| pick(b, a))?;
if let Some(q) = outer_col.qualifier.as_deref()
&& !q.eq_ignore_ascii_case(outer_alias)
{
return None;
}
let outer_pos = outer_schema
.iter()
.position(|c| c.name.eq_ignore_ascii_case(&outer_col.name))?;
let catalog = self.active_catalog();
let table_idx = catalog.tables_position_of(inner_table_name.as_str())?;
let inner_table = catalog.tables_at(table_idx)?;
let inner_schema_ref = inner_table.schema();
let inner_pos = inner_schema_ref
.columns
.iter()
.position(|c| c.name.eq_ignore_ascii_case(&inner_col_name))?;
if !matches!(
inner_schema_ref.columns[inner_pos].ty,
spg_storage::DataType::BigInt
| spg_storage::DataType::Int
| spg_storage::DataType::SmallInt
) {
return None;
}
if !inner_schema_ref
.uniqueness_constraints
.iter()
.any(|u| u.is_primary_key && u.columns.as_slice() == [inner_pos])
{
return None;
}
Some(ScalarPkProbeFastPath {
outer_pos,
inner_table_name,
inner_pos,
table_idx,
})
}
pub(crate) fn try_scalar_count_pk_eq_probe(
&self,
inner: &SelectStatement,
row: &Row<'static>,
ctx: &EvalContext<'_>,
) -> Result<Option<Value<'static>>, EngineError> {
use spg_sql::ast::{BinOp, ColumnName, SelectItem};
if !inner.ctes.is_empty()
|| !inner.unions.is_empty()
|| inner.group_by.is_some()
|| inner.having.is_some()
|| inner.distinct
|| !inner.order_by.is_empty()
|| inner.limit.is_some()
|| inner.offset.is_some()
|| inner.items.len() != 1
{
return Ok(None);
}
let SelectItem::Expr { expr, .. } = &inner.items[0] else {
return Ok(None);
};
let is_count_shape = match expr {
Expr::FunctionCall { name, args } => {
(name.eq_ignore_ascii_case("count_star") && args.is_empty())
|| name.eq_ignore_ascii_case("count")
}
_ => false,
};
if !is_count_shape {
return Ok(None);
}
let Some(from) = &inner.from else {
return Ok(None);
};
if !from.joins.is_empty()
|| 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 inner_table_name = from.primary.name.as_str();
let inner_alias = from.primary.alias.as_deref().unwrap_or(inner_table_name);
let Some(where_expr) = &inner.where_ else {
return Ok(None);
};
let Expr::Binary {
lhs,
op: BinOp::Eq,
rhs,
} = where_expr
else {
return Ok(None);
};
let (Expr::Column(a), Expr::Column(b)) = (lhs.as_ref(), rhs.as_ref()) else {
return Ok(None);
};
let pick = |x: &ColumnName, y: &ColumnName| -> Option<(String, ColumnName)> {
if x.qualifier
.as_deref()
.is_some_and(|q| q.eq_ignore_ascii_case(inner_alias))
{
Some((x.name.clone(), y.clone()))
} else {
None
}
};
let Some((inner_col_name, outer_col)) = pick(a, b).or_else(|| pick(b, a)) else {
return Ok(None);
};
let catalog = self.active_catalog();
let Some(inner_table) = catalog.get(inner_table_name) else {
return Ok(None);
};
let inner_schema = inner_table.schema();
let Some(inner_pos) = inner_schema
.columns
.iter()
.position(|c| c.name.eq_ignore_ascii_case(&inner_col_name))
else {
return Ok(None);
};
if !matches!(
inner_schema.columns[inner_pos].ty,
spg_storage::DataType::BigInt
| spg_storage::DataType::Int
| spg_storage::DataType::SmallInt
) {
return Ok(None);
}
if !inner_schema
.uniqueness_constraints
.iter()
.any(|u| u.is_primary_key && u.columns.as_slice() == [inner_pos])
{
return Ok(None);
}
let outer_val = match eval::eval_expr(&Expr::Column(outer_col), row, ctx) {
Ok(v) => v,
Err(_) => return Ok(None),
};
let outer_int = match outer_val {
Value::BigInt(n) => n,
Value::Int(n) => i64::from(n),
Value::SmallInt(n) => i64::from(n),
Value::Null => return Ok(Some(Value::BigInt(0))),
_ => return Ok(None),
};
let Some(idx) = inner_table.index_on(inner_pos) else {
return Ok(None);
};
let Some(key) = spg_storage::IndexKey::from_value(&Value::BigInt(outer_int)) else {
return Ok(None);
};
SCALARSQ_PK_PROBE_FIRED.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
let hit = !idx.lookup_eq(&key).is_empty();
Ok(Some(Value::BigInt(i64::from(hit))))
}
}
pub static SCALARSQ_PK_PROBE_FIRED: core::sync::atomic::AtomicU64 =
core::sync::atomic::AtomicU64::new(0);
fn scalar_subquery_empty_default(inner: &SelectStatement) -> Value<'static> {
use spg_sql::ast::SelectItem;
if inner.items.len() != 1 {
return Value::Null;
}
let SelectItem::Expr { expr, .. } = &inner.items[0] else {
return Value::Null;
};
fn is_count(e: &Expr) -> bool {
match e {
Expr::FunctionCall { name, .. } => {
name.eq_ignore_ascii_case("count") || name.eq_ignore_ascii_case("count_star")
}
Expr::AggregateOrdered { call, .. } => is_count(call),
_ => false,
}
}
if is_count(expr) {
Value::BigInt(0)
} else {
Value::Null
}
}
fn relation_name_matches(qualifier: &str, relation: &str) -> bool {
if qualifier.eq_ignore_ascii_case(relation) {
return true;
}
let rewritten = if let Some(bare) = qualifier
.to_ascii_lowercase()
.strip_prefix("pg_")
.map(alloc::string::String::from)
{
alloc::format!("__spg_pg_{bare}")
} else {
alloc::format!("__spg_info_{}", qualifier.to_ascii_lowercase())
};
rewritten.eq_ignore_ascii_case(relation)
}
fn inner_scope_column_names(
s: &SelectStatement,
cat: &spg_storage::Catalog,
) -> Option<alloc::collections::BTreeSet<alloc::string::String>> {
use spg_sql::ast::SelectItem;
fn add_table(
t: &spg_sql::ast::TableRef,
cat: &spg_storage::Catalog,
names: &mut alloc::collections::BTreeSet<alloc::string::String>,
) -> bool {
if let Some(body) = &t.lateral_subquery {
if body.items.is_empty() {
return false;
}
for item in &body.items {
match item {
SelectItem::Expr { alias: Some(a), .. } => {
names.insert(a.to_ascii_lowercase());
}
SelectItem::Expr {
expr: Expr::Column(c),
..
} => {
names.insert(c.name.to_ascii_lowercase());
}
SelectItem::Wildcard | SelectItem::QualifiedWildcard(_) => {
let Some(inner) = inner_scope_column_names(body, cat) else {
return false;
};
names.extend(inner);
}
SelectItem::Expr { .. } => return false,
}
}
return true;
}
if t.unnest_expr.is_some() || t.generate_series_args.is_some() || t.name.is_empty() {
return false;
}
let Some(tbl) = cat.get(&t.name) else {
return false;
};
for col in &tbl.schema().columns {
names.insert(col.name.to_ascii_lowercase());
}
true
}
let mut names: alloc::collections::BTreeSet<alloc::string::String> =
alloc::collections::BTreeSet::new();
let from = s.from.as_ref()?;
if !add_table(&from.primary, cat, &mut names) {
return None;
}
for j in &from.joins {
if !add_table(&j.table, cat, &mut names) {
return None;
}
}
for item in &s.items {
if let SelectItem::Expr { alias: Some(a), .. } = item {
names.insert(a.to_ascii_lowercase());
}
}
if !s.ctes.is_empty() {
return None;
}
Some(names)
}
fn is_synthetic_column_name(n: &str) -> bool {
n.starts_with("__grp_") || n.starts_with("__agg_") || n.starts_with("__spg_")
}
pub(crate) fn select_is_correlated(s: &SelectStatement) -> bool {
use spg_sql::ast::SelectItem;
let Some(from) = &s.from else {
let mut qualified = false;
for item in &s.items {
if let SelectItem::Expr { expr, .. } = item {
visit_expr_columns_and_subqueries(
expr,
&mut |c| {
if c.qualifier.is_some() {
qualified = true;
}
},
&mut |_| {},
);
}
}
return qualified;
};
let mut inner: Vec<&str> = Vec::new();
if let Some(a) = &from.primary.alias {
inner.push(a.as_str());
}
if !from.primary.name.is_empty() {
inner.push(from.primary.name.as_str());
}
for j in &from.joins {
if let Some(a) = &j.table.alias {
inner.push(a.as_str());
}
if !j.table.name.is_empty() {
inner.push(j.table.name.as_str());
}
}
let mut exprs: Vec<&Expr> = Vec::new();
for item in &s.items {
if let SelectItem::Expr { expr, .. } = item {
exprs.push(expr);
}
}
if let Some(w) = &s.where_ {
exprs.push(w);
}
for j in &from.joins {
if let Some(on) = &j.on {
exprs.push(on);
}
}
if let Some(gs) = &s.group_by {
for g in gs {
exprs.push(g);
}
}
if let Some(h) = &s.having {
exprs.push(h);
}
for o in &s.order_by {
exprs.push(&o.expr);
}
let mut correlated = false;
for e in exprs {
visit_expr_columns_and_subqueries(
e,
&mut |c| {
if let Some(q) = &c.qualifier
&& !inner.iter().any(|a| relation_name_matches(q, a))
{
correlated = true;
}
},
&mut |_| {},
);
}
if !correlated {
for t in core::iter::once(&from.primary).chain(from.joins.iter().map(|j| &j.table)) {
if let Some(body) = &t.lateral_subquery
&& select_is_correlated(body)
{
correlated = true;
break;
}
}
}
correlated
}
pub(crate) fn collect_scalar_subqueries<'a>(e: &'a Expr, out: &mut Vec<&'a SelectStatement>) {
match e {
Expr::ScalarSubquery(s) => out.push(s),
Expr::Exists { .. }
| Expr::InSubquery { .. }
| Expr::RowInSubquery { .. }
| Expr::RowCmpSubquery { .. } => {}
Expr::Binary { lhs, rhs, .. } => {
collect_scalar_subqueries(lhs, out);
collect_scalar_subqueries(rhs, out);
}
Expr::Unary { expr, .. }
| Expr::Cast { expr, .. }
| Expr::IsNull { expr, .. }
| Expr::BoolTest { expr, .. }
| Expr::FieldAccess { base: expr, .. } => {
collect_scalar_subqueries(expr, out);
}
Expr::Like { expr, pattern, .. } => {
collect_scalar_subqueries(expr, out);
collect_scalar_subqueries(pattern, out);
}
Expr::FunctionCall { args, .. } => {
for a in args {
collect_scalar_subqueries(a, out);
}
}
Expr::AggregateOrdered { call, order_by, .. } => {
collect_scalar_subqueries(call, out);
for o in order_by {
collect_scalar_subqueries(&o.expr, out);
}
}
Expr::Case {
operand,
branches,
else_branch,
} => {
if let Some(op) = operand {
collect_scalar_subqueries(op, out);
}
for (w, t) in branches {
collect_scalar_subqueries(w, out);
collect_scalar_subqueries(t, out);
}
if let Some(eb) = else_branch {
collect_scalar_subqueries(eb, out);
}
}
Expr::ArraySubscript { target, index } => {
collect_scalar_subqueries(target, out);
collect_scalar_subqueries(index, out);
}
Expr::InList { expr, list, .. } => {
collect_scalar_subqueries(expr, out);
for item in list {
collect_scalar_subqueries(item, out);
}
}
_ => {}
}
}
fn hollow_scalar_subqueries(e: &mut Expr) {
match e {
Expr::ScalarSubquery(s) => {
let hollow = SelectStatement {
items: Vec::new(),
..SelectStatement::default()
};
**s = hollow;
}
Expr::Exists { .. }
| Expr::InSubquery { .. }
| Expr::RowInSubquery { .. }
| Expr::RowCmpSubquery { .. } => {}
Expr::Binary { lhs, rhs, .. } => {
hollow_scalar_subqueries(lhs);
hollow_scalar_subqueries(rhs);
}
Expr::Unary { expr, .. }
| Expr::Cast { expr, .. }
| Expr::IsNull { expr, .. }
| Expr::BoolTest { expr, .. }
| Expr::FieldAccess { base: expr, .. } => {
hollow_scalar_subqueries(expr);
}
Expr::Like { expr, pattern, .. } => {
hollow_scalar_subqueries(expr);
hollow_scalar_subqueries(pattern);
}
Expr::FunctionCall { args, .. } => {
for a in args.iter_mut() {
hollow_scalar_subqueries(a);
}
}
Expr::AggregateOrdered { call, order_by, .. } => {
hollow_scalar_subqueries(call);
for o in order_by.iter_mut() {
hollow_scalar_subqueries(&mut o.expr);
}
}
Expr::Case {
operand,
branches,
else_branch,
} => {
if let Some(op) = operand {
hollow_scalar_subqueries(op);
}
for (w, t) in branches.iter_mut() {
hollow_scalar_subqueries(w);
hollow_scalar_subqueries(t);
}
if let Some(eb) = else_branch {
hollow_scalar_subqueries(eb);
}
}
Expr::ArraySubscript { target, index } => {
hollow_scalar_subqueries(target);
hollow_scalar_subqueries(index);
}
Expr::InList { expr, list, .. } => {
hollow_scalar_subqueries(expr);
for item in list.iter_mut() {
hollow_scalar_subqueries(item);
}
}
_ => {}
}
}
fn splice_planned_subqueries(
e: &mut Expr,
plan: &[Option<alloc::rc::Rc<memoize::GroupMap>>],
idx: &mut usize,
row: &Row<'static>,
ctx: &EvalContext<'_>,
) -> Result<bool, EngineError> {
match e {
Expr::ScalarSubquery(_) => {
let Some(Some(gm)) = plan.get(*idx) else {
return Ok(false);
};
*idx += 1;
let (outer_col, map, empty_default) = gm.as_ref();
let key_v = eval::eval_expr(&Expr::Column(outer_col.clone()), row, ctx)
.map_err(EngineError::Eval)?;
let v = map
.get(&aggregate::encode_key(core::slice::from_ref(&key_v)))
.cloned()
.unwrap_or_else(|| empty_default.clone());
*e = value_to_literal_expr(v)?;
Ok(true)
}
Expr::Exists { .. }
| Expr::InSubquery { .. }
| Expr::RowInSubquery { .. }
| Expr::RowCmpSubquery { .. } => Ok(true),
Expr::Binary { lhs, rhs, .. } => Ok(splice_planned_subqueries(lhs, plan, idx, row, ctx)?
&& splice_planned_subqueries(rhs, plan, idx, row, ctx)?),
Expr::Unary { expr, .. }
| Expr::Cast { expr, .. }
| Expr::IsNull { expr, .. }
| Expr::BoolTest { expr, .. }
| Expr::FieldAccess { base: expr, .. } => {
splice_planned_subqueries(expr, plan, idx, row, ctx)
}
Expr::Like { expr, pattern, .. } => {
Ok(splice_planned_subqueries(expr, plan, idx, row, ctx)?
&& splice_planned_subqueries(pattern, plan, idx, row, ctx)?)
}
Expr::FunctionCall { args, .. } => {
for a in args.iter_mut() {
if !splice_planned_subqueries(a, plan, idx, row, ctx)? {
return Ok(false);
}
}
Ok(true)
}
Expr::AggregateOrdered { call, order_by, .. } => {
if !splice_planned_subqueries(call, plan, idx, row, ctx)? {
return Ok(false);
}
for o in order_by.iter_mut() {
if !splice_planned_subqueries(&mut o.expr, plan, idx, row, ctx)? {
return Ok(false);
}
}
Ok(true)
}
Expr::Case {
operand,
branches,
else_branch,
} => {
if let Some(op) = operand {
if !splice_planned_subqueries(op, plan, idx, row, ctx)? {
return Ok(false);
}
}
for (w, t) in branches.iter_mut() {
if !splice_planned_subqueries(w, plan, idx, row, ctx)?
|| !splice_planned_subqueries(t, plan, idx, row, ctx)?
{
return Ok(false);
}
}
if let Some(eb) = else_branch {
if !splice_planned_subqueries(eb, plan, idx, row, ctx)? {
return Ok(false);
}
}
Ok(true)
}
Expr::ArraySubscript { target, index } => {
Ok(splice_planned_subqueries(target, plan, idx, row, ctx)?
&& splice_planned_subqueries(index, plan, idx, row, ctx)?)
}
Expr::InList { expr, list, .. } => {
if !splice_planned_subqueries(expr, plan, idx, row, ctx)? {
return Ok(false);
}
for item in list.iter_mut() {
if !splice_planned_subqueries(item, plan, idx, row, ctx)? {
return Ok(false);
}
}
Ok(true)
}
_ => Ok(true),
}
}
pub(crate) fn collect_exists_subqueries<'a>(e: &'a Expr, out: &mut Vec<&'a SelectStatement>) {
match e {
Expr::Exists { subquery, .. } => out.push(subquery.as_ref()),
Expr::ScalarSubquery(_)
| Expr::InSubquery { .. }
| Expr::RowInSubquery { .. }
| Expr::RowCmpSubquery { .. } => {}
Expr::Binary { lhs, rhs, .. } => {
collect_exists_subqueries(lhs, out);
collect_exists_subqueries(rhs, out);
}
Expr::Unary { expr, .. }
| Expr::Cast { expr, .. }
| Expr::IsNull { expr, .. }
| Expr::BoolTest { expr, .. }
| Expr::FieldAccess { base: expr, .. } => {
collect_exists_subqueries(expr, out);
}
Expr::Like { expr, pattern, .. } => {
collect_exists_subqueries(expr, out);
collect_exists_subqueries(pattern, out);
}
Expr::FunctionCall { args, .. } => {
for a in args {
collect_exists_subqueries(a, out);
}
}
Expr::AggregateOrdered { call, order_by, .. } => {
collect_exists_subqueries(call, out);
for o in order_by {
collect_exists_subqueries(&o.expr, out);
}
}
Expr::Case {
operand,
branches,
else_branch,
} => {
if let Some(op) = operand {
collect_exists_subqueries(op, out);
}
for (w, t) in branches {
collect_exists_subqueries(w, out);
collect_exists_subqueries(t, out);
}
if let Some(eb) = else_branch {
collect_exists_subqueries(eb, out);
}
}
Expr::ArraySubscript { target, index } => {
collect_exists_subqueries(target, out);
collect_exists_subqueries(index, out);
}
Expr::InList { expr, list, .. } => {
collect_exists_subqueries(expr, out);
for item in list {
collect_exists_subqueries(item, out);
}
}
_ => {}
}
}
fn bare_exists_shape(e: &Expr) -> Option<(bool, bool)> {
match e {
Expr::Exists { negated, .. } => Some((*negated, false)),
Expr::Unary {
op: spg_sql::ast::UnOp::Not,
expr: inner,
} => match inner.as_ref() {
Expr::Exists { negated, .. } => Some((*negated, true)),
_ => None,
},
_ => None,
}
}
fn planned_exists_bit(
es: &memoize::ExistsSet,
negated: bool,
row: &Row<'static>,
ctx: &EvalContext<'_>,
) -> Result<bool, EngineError> {
let (outer_cols, set) = es;
let mut key_vals: Vec<Value<'static>> = Vec::with_capacity(outer_cols.len());
let mut any_null = false;
for oc in outer_cols {
let v = eval::eval_expr(oc, row, ctx).map_err(EngineError::Eval)?;
if matches!(v, Value::Null) {
any_null = true;
}
key_vals.push(v);
}
let present = !any_null && set.contains(&aggregate::encode_canonical_key(&key_vals));
Ok(if negated { !present } else { present })
}
fn splice_planned_exists(
e: &mut Expr,
plan: &[Option<alloc::rc::Rc<memoize::ExistsSet>>],
idx: &mut usize,
row: &Row<'static>,
ctx: &EvalContext<'_>,
) -> Result<bool, EngineError> {
match e {
Expr::Exists { negated, .. } => {
let Some(Some(es)) = plan.get(*idx) else {
return Ok(false);
};
*idx += 1;
let bit = planned_exists_bit(es, *negated, row, ctx)?;
*e = Expr::Literal(Literal::Bool(bit));
Ok(true)
}
Expr::ScalarSubquery(_)
| Expr::InSubquery { .. }
| Expr::RowInSubquery { .. }
| Expr::RowCmpSubquery { .. } => Ok(true),
Expr::Binary { lhs, rhs, .. } => Ok(splice_planned_exists(lhs, plan, idx, row, ctx)?
&& splice_planned_exists(rhs, plan, idx, row, ctx)?),
Expr::Unary { expr, .. }
| Expr::Cast { expr, .. }
| Expr::IsNull { expr, .. }
| Expr::BoolTest { expr, .. }
| Expr::FieldAccess { base: expr, .. } => splice_planned_exists(expr, plan, idx, row, ctx),
Expr::Like { expr, pattern, .. } => Ok(splice_planned_exists(expr, plan, idx, row, ctx)?
&& splice_planned_exists(pattern, plan, idx, row, ctx)?),
Expr::FunctionCall { args, .. } => {
for a in args.iter_mut() {
if !splice_planned_exists(a, plan, idx, row, ctx)? {
return Ok(false);
}
}
Ok(true)
}
Expr::AggregateOrdered { call, order_by, .. } => {
if !splice_planned_exists(call, plan, idx, row, ctx)? {
return Ok(false);
}
for o in order_by.iter_mut() {
if !splice_planned_exists(&mut o.expr, plan, idx, row, ctx)? {
return Ok(false);
}
}
Ok(true)
}
Expr::Case {
operand,
branches,
else_branch,
} => {
if let Some(op) = operand {
if !splice_planned_exists(op, plan, idx, row, ctx)? {
return Ok(false);
}
}
for (w, t) in branches.iter_mut() {
if !splice_planned_exists(w, plan, idx, row, ctx)?
|| !splice_planned_exists(t, plan, idx, row, ctx)?
{
return Ok(false);
}
}
if let Some(eb) = else_branch {
if !splice_planned_exists(eb, plan, idx, row, ctx)? {
return Ok(false);
}
}
Ok(true)
}
Expr::ArraySubscript { target, index } => {
Ok(splice_planned_exists(target, plan, idx, row, ctx)?
&& splice_planned_exists(index, plan, idx, row, ctx)?)
}
Expr::InList { expr, list, .. } => {
if !splice_planned_exists(expr, plan, idx, row, ctx)? {
return Ok(false);
}
for item in list.iter_mut() {
if !splice_planned_exists(item, plan, idx, row, ctx)? {
return Ok(false);
}
}
Ok(true)
}
_ => Ok(true),
}
}
const INLIST_SET_THRESHOLD: usize = 64;
fn expr_may_use_in_set(e: &Expr) -> bool {
match e {
Expr::InList { list, .. } => list.len() >= INLIST_SET_THRESHOLD,
Expr::Binary {
lhs,
op: BinOp::And,
rhs,
} => expr_may_use_in_set(lhs) || expr_may_use_in_set(rhs),
_ => false,
}
}
fn cast_target_is_integer(target: &spg_sql::ast::CastTarget) -> bool {
use spg_sql::ast::CastTarget;
match target {
CastTarget::BigInt | CastTarget::Int => true,
CastTarget::Named(n) => {
matches!(
n.to_ascii_lowercase().as_str(),
"int2" | "int4" | "int8" | "smallint" | "integer" | "int" | "bigint"
)
}
_ => false,
}
}
pub(crate) fn build_in_list_set(list: &[Expr]) -> Option<memoize::InListSetEntry> {
let mut has_null = false;
let mut ints: hashbrown::HashSet<i64> = hashbrown::HashSet::with_capacity(list.len());
let mut texts: hashbrown::HashSet<String> = hashbrown::HashSet::with_capacity(list.len());
for item in list {
let lit = match item {
Expr::Literal(lit) => lit,
Expr::Cast { expr, target } if cast_target_is_integer(target) => match expr.as_ref() {
Expr::Literal(inner) => inner,
_ => return None,
},
_ => return None,
};
match lit {
Literal::Null => has_null = true,
Literal::Integer(i) => {
ints.insert(*i);
}
Literal::String(s) => {
texts.insert(s.clone());
}
_ => return None,
}
if !ints.is_empty() && !texts.is_empty() {
return None;
}
}
let set = if !ints.is_empty() {
memoize::InListSet::Int(ints)
} else if !texts.is_empty() {
memoize::InListSet::Text(texts)
} else {
return None;
};
Some(memoize::InListSetEntry { set, has_null })
}
fn eval_with_in_sets(
e: &Expr,
row: &Row<'static>,
ctx: &EvalContext<'_>,
m: &mut memoize::MemoizeCache,
) -> Result<Value<'static>, EngineError> {
match e {
Expr::Binary {
lhs,
op: BinOp::And,
rhs,
} => {
let l = eval_with_in_sets(lhs, row, ctx, m)?;
let r = eval_with_in_sets(rhs, row, ctx, m)?;
eval::and_3vl(l, r).map_err(EngineError::Eval)
}
Expr::InList {
expr: lhs,
list,
negated,
} if list.len() >= INLIST_SET_THRESHOLD => {
let key = core::ptr::from_ref::<Expr>(e) as usize;
let Some(entry) = m
.in_sets
.entry(key)
.or_insert_with(|| build_in_list_set(list))
else {
return eval::eval_expr(e, row, ctx).map_err(EngineError::Eval);
};
let needle = eval::eval_expr(lhs, row, ctx).map_err(EngineError::Eval)?;
let contained = match (&needle, &entry.set) {
(Value::Null, _) => return Ok(Value::Null),
(Value::SmallInt(n), memoize::InListSet::Int(s)) => s.contains(&i64::from(*n)),
(Value::Int(n), memoize::InListSet::Int(s)) => s.contains(&i64::from(*n)),
(Value::BigInt(n), memoize::InListSet::Int(s)) => s.contains(n),
(Value::Text(t), memoize::InListSet::Text(s)) => s.contains(t.as_ref()),
_ => return eval::eval_expr(e, row, ctx).map_err(EngineError::Eval),
};
let inner = if contained {
Value::Bool(true)
} else if entry.has_null {
Value::Null
} else {
Value::Bool(false)
};
Ok(match (negated, inner) {
(true, Value::Bool(b)) => Value::Bool(!b),
(_, v) => v,
})
}
_ => eval::eval_expr(e, row, ctx).map_err(EngineError::Eval),
}
}
fn substitute_outer_columns(
stmt: &mut SelectStatement,
row: &Row<'static>,
ctx: &EvalContext<'_>,
cat: &spg_storage::Catalog,
) {
let outer_alias = ctx.table_alias.unwrap_or("");
substitute_in_select(stmt, row, ctx, outer_alias, cat);
}
fn substitute_in_select(
stmt: &mut SelectStatement,
row: &Row<'static>,
ctx: &EvalContext<'_>,
outer_alias: &str,
cat: &spg_storage::Catalog,
) {
let visible = inner_scope_column_names(stmt, cat);
for item in &mut stmt.items {
if let SelectItem::Expr { expr, .. } = item {
substitute_in_expr(expr, row, ctx, outer_alias, cat, visible.as_ref());
}
}
if let Some(w) = &mut stmt.where_ {
substitute_in_expr(w, row, ctx, outer_alias, cat, visible.as_ref());
}
if let Some(gs) = &mut stmt.group_by {
for g in gs {
substitute_in_expr(g, row, ctx, outer_alias, cat, visible.as_ref());
}
}
if let Some(h) = &mut stmt.having {
substitute_in_expr(h, row, ctx, outer_alias, cat, visible.as_ref());
}
for o in &mut stmt.order_by {
substitute_in_expr(&mut o.expr, row, ctx, outer_alias, cat, visible.as_ref());
}
for (_, peer) in &mut stmt.unions {
substitute_in_select(peer, row, ctx, outer_alias, cat);
}
if let Some(from) = &mut stmt.from {
if let Some(body) = &mut from.primary.lateral_subquery {
substitute_in_select(body, row, ctx, outer_alias, cat);
}
for j in &mut from.joins {
if let Some(on) = &mut j.on {
substitute_in_expr(on, row, ctx, outer_alias, cat, visible.as_ref());
}
if let Some(body) = &mut j.table.lateral_subquery {
substitute_in_select(body, row, ctx, outer_alias, cat);
}
}
}
}
fn substitute_in_expr(
e: &mut Expr,
row: &Row<'static>,
ctx: &EvalContext<'_>,
outer_alias: &str,
cat: &spg_storage::Catalog,
visible: Option<&alloc::collections::BTreeSet<alloc::string::String>>,
) {
if let Expr::Column(c) = e
&& c.qualifier.is_none()
&& (c.name.starts_with("__grp_") || c.name.starts_with("__agg_"))
&& let Some(idx) = ctx.columns.iter().position(|sc| sc.name == c.name)
{
let v = row.values.get(idx).cloned().unwrap_or(Value::Null);
if let Ok(lit) = value_to_literal_expr(v) {
*e = lit;
return;
}
}
if let Expr::Column(c) = e
&& c.qualifier.is_none()
&& c.name != "*"
&& !is_synthetic_column_name(&c.name.to_ascii_lowercase())
&& visible.is_some_and(|v| !v.contains(&c.name.to_ascii_lowercase()))
&& let Some(idx) = ctx
.columns
.iter()
.position(|sc| sc.name.eq_ignore_ascii_case(&c.name))
{
let v = row.values.get(idx).cloned().unwrap_or(Value::Null);
if let Ok(lit) = value_to_literal_expr(v) {
*e = lit;
return;
}
}
if let Expr::Column(c) = e
&& let Some(qual) = &c.qualifier
{
let idx = if !outer_alias.is_empty() && relation_name_matches(qual, outer_alias) {
ctx.columns
.iter()
.position(|sc| sc.name.eq_ignore_ascii_case(&c.name))
} else {
None
}
.or_else(|| {
let composite = alloc::format!("{qual}.{name}", name = c.name);
ctx.columns
.iter()
.position(|sc| sc.name.eq_ignore_ascii_case(&composite))
});
if let Some(idx) = idx {
let v = row.values.get(idx).cloned().unwrap_or(Value::Null);
if let Ok(lit) = value_to_literal_expr(v) {
*e = lit;
return;
}
}
}
match e {
Expr::NamedArg { expr, .. } => {
substitute_in_expr(expr, row, ctx, outer_alias, cat, visible)
}
Expr::Variadic(expr) => substitute_in_expr(expr, row, ctx, outer_alias, cat, visible),
Expr::AggregateOrdered { call, order_by, .. } => {
substitute_in_expr(call, row, ctx, outer_alias, cat, visible);
for o in order_by.iter_mut() {
substitute_in_expr(&mut o.expr, row, ctx, outer_alias, cat, visible);
}
}
Expr::Binary { lhs, rhs, .. } => {
substitute_in_expr(lhs, row, ctx, outer_alias, cat, visible);
substitute_in_expr(rhs, row, ctx, outer_alias, cat, visible);
}
Expr::Unary { expr, .. }
| Expr::Cast { expr, .. }
| Expr::IsNull { expr, .. }
| Expr::BoolTest { expr, .. }
| Expr::FieldAccess { base: expr, .. } => {
substitute_in_expr(expr, row, ctx, outer_alias, cat, visible);
}
Expr::Like { expr, pattern, .. } => {
substitute_in_expr(expr, row, ctx, outer_alias, cat, visible);
substitute_in_expr(pattern, row, ctx, outer_alias, cat, visible);
}
Expr::FunctionCall { args, .. } => {
for a in args {
substitute_in_expr(a, row, ctx, outer_alias, cat, visible);
}
}
Expr::Extract { source, .. } => {
substitute_in_expr(source, row, ctx, outer_alias, cat, visible)
}
Expr::WindowFunction {
args,
partition_by,
order_by,
..
} => {
for a in args {
substitute_in_expr(a, row, ctx, outer_alias, cat, visible);
}
for p in partition_by {
substitute_in_expr(p, row, ctx, outer_alias, cat, visible);
}
for (o, _, _) in order_by {
substitute_in_expr(o, row, ctx, outer_alias, cat, visible);
}
}
Expr::ScalarSubquery(s) => substitute_in_select(s, row, ctx, outer_alias, cat),
Expr::Exists { subquery, .. } | Expr::InSubquery { subquery, .. } => {
substitute_in_select(subquery, row, ctx, outer_alias, cat);
}
Expr::RowInSubquery {
row: row_exprs,
subquery,
..
} => {
for el in row_exprs.iter_mut() {
substitute_in_expr(el, row, ctx, outer_alias, cat, visible);
}
substitute_in_select(subquery, row, ctx, outer_alias, cat);
}
Expr::RowCmpSubquery {
row: row_exprs,
subquery,
..
} => {
for el in row_exprs.iter_mut() {
substitute_in_expr(el, row, ctx, outer_alias, cat, visible);
}
substitute_in_select(subquery, row, ctx, outer_alias, cat);
}
Expr::Literal(_) | Expr::Placeholder(_) | Expr::Column(_) => {}
Expr::Array(items) => {
for elem in items {
substitute_in_expr(elem, row, ctx, outer_alias, cat, visible);
}
}
Expr::ArraySubscript { target, index } => {
substitute_in_expr(target, row, ctx, outer_alias, cat, visible);
substitute_in_expr(index, row, ctx, outer_alias, cat, visible);
}
Expr::ArraySlice { target, lo, hi } => {
substitute_in_expr(target, row, ctx, outer_alias, cat, visible);
if let Some(l) = lo {
substitute_in_expr(l, row, ctx, outer_alias, cat, visible);
}
if let Some(h) = hi {
substitute_in_expr(h, row, ctx, outer_alias, cat, visible);
}
}
Expr::AnyAll { expr, array, .. } => {
substitute_in_expr(expr, row, ctx, outer_alias, cat, visible);
substitute_in_expr(array, row, ctx, outer_alias, cat, visible);
}
Expr::InList { expr, list, .. } => {
substitute_in_expr(expr, row, ctx, outer_alias, cat, visible);
for item in list {
substitute_in_expr(item, row, ctx, outer_alias, cat, visible);
}
}
Expr::Case {
operand,
branches,
else_branch,
} => {
if let Some(o) = operand {
substitute_in_expr(o, row, ctx, outer_alias, cat, visible);
}
for (w, t) in branches {
substitute_in_expr(w, row, ctx, outer_alias, cat, visible);
substitute_in_expr(t, row, ctx, outer_alias, cat, visible);
}
if let Some(e) = else_branch {
substitute_in_expr(e, row, ctx, outer_alias, cat, visible);
}
}
}
}
pub fn expr_tree_has_subquery(stmt: &SelectStatement) -> bool {
let mut any = false;
for item in &stmt.items {
if let SelectItem::Expr { expr, .. } = item {
any = any || expr_has_subquery(expr);
}
}
if let Some(w) = &stmt.where_ {
any = any || expr_has_subquery(w);
}
if let Some(h) = &stmt.having {
any = any || expr_has_subquery(h);
}
for o in &stmt.order_by {
any = any || expr_has_subquery(&o.expr);
}
for (_, peer) in &stmt.unions {
any = any || expr_tree_has_subquery(peer);
}
any
}
pub(crate) fn expr_has_subquery(e: &Expr) -> bool {
match e {
Expr::NamedArg { expr, .. } => expr_has_subquery(expr),
Expr::Variadic(expr) => expr_has_subquery(expr),
Expr::ScalarSubquery(_)
| Expr::Exists { .. }
| Expr::InSubquery { .. }
| Expr::RowInSubquery { .. }
| Expr::RowCmpSubquery { .. } => true,
Expr::AggregateOrdered { call, order_by, .. } => {
expr_has_subquery(call) || order_by.iter().any(|o| expr_has_subquery(&o.expr))
}
Expr::Binary { lhs, rhs, .. } => expr_has_subquery(lhs) || expr_has_subquery(rhs),
Expr::Unary { expr, .. }
| Expr::Cast { expr, .. }
| Expr::IsNull { expr, .. }
| Expr::BoolTest { expr, .. }
| Expr::FieldAccess { base: expr, .. } => expr_has_subquery(expr),
Expr::FunctionCall { args, .. } => args.iter().any(expr_has_subquery),
Expr::Like { expr, pattern, .. } => expr_has_subquery(expr) || expr_has_subquery(pattern),
Expr::Extract { source, .. } => expr_has_subquery(source),
Expr::WindowFunction {
args,
partition_by,
order_by,
..
} => {
args.iter().any(expr_has_subquery)
|| partition_by.iter().any(expr_has_subquery)
|| order_by.iter().any(|(e, _, _)| expr_has_subquery(e))
}
Expr::Literal(_) | Expr::Placeholder(_) | Expr::Column(_) => false,
Expr::Array(items) => items.iter().any(expr_has_subquery),
Expr::ArraySubscript { target, index } => {
expr_has_subquery(target) || expr_has_subquery(index)
}
Expr::ArraySlice { target, lo, hi } => {
expr_has_subquery(target)
|| lo.as_deref().is_some_and(expr_has_subquery)
|| hi.as_deref().is_some_and(expr_has_subquery)
}
Expr::AnyAll { expr, array, .. } => expr_has_subquery(expr) || expr_has_subquery(array),
Expr::InList { expr, list, .. } => {
expr_has_subquery(expr) || list.iter().any(expr_has_subquery)
}
Expr::Case {
operand,
branches,
else_branch,
} => {
operand.as_deref().is_some_and(expr_has_subquery)
|| branches
.iter()
.any(|(w, t)| expr_has_subquery(w) || expr_has_subquery(t))
|| else_branch.as_deref().is_some_and(expr_has_subquery)
}
}
}