1use alopex_core::kv::KVStore;
2use alopex_core::sql::subquery::{materialize_cache, nested_scan, semi_join_probe};
3
4use crate::catalog::Catalog;
5use crate::executor::evaluator::EvalContext;
6use crate::executor::{EvaluationError, ExecutorError, Result, Row};
7use crate::planner::logical_plan::LogicalPlan;
8use crate::planner::typed_expr::{Quantifier, TypedExpr, TypedExprKind};
9use crate::storage::{SqlTxn, SqlValue};
10
11pub fn execute_scalar_subquery<'txn, S: KVStore + 'txn, C: Catalog + ?Sized, T: SqlTxn<'txn, S>>(
13 txn: &mut T,
14 catalog: &C,
15 subquery: &LogicalPlan,
16) -> Result<SqlValue> {
17 execute_scalar_subquery_with_outer(txn, catalog, subquery, None)
18}
19
20pub(crate) fn execute_scalar_subquery_with_outer<
21 'txn,
22 S: KVStore + 'txn,
23 C: Catalog + ?Sized,
24 T: SqlTxn<'txn, S>,
25>(
26 txn: &mut T,
27 catalog: &C,
28 subquery: &LogicalPlan,
29 outer: Option<&Row>,
30) -> Result<SqlValue> {
31 let rows = execute_subquery_rows_with_outer(txn, catalog, subquery, outer)?;
32 if rows.len() > 1 {
33 return Err(ExecutorError::InvalidOperation {
34 operation: "execute_scalar_subquery".into(),
35 reason: "scalar subquery returned multiple rows".into(),
36 });
37 }
38 let Some(row) = rows.first() else {
39 return Ok(SqlValue::Null);
40 };
41 if row.len() != 1 {
42 return Err(ExecutorError::InvalidOperation {
43 operation: "execute_scalar_subquery".into(),
44 reason: format!("scalar subquery returned {} columns", row.len()),
45 });
46 }
47 Ok(row[0].clone())
48}
49
50pub fn execute_in_subquery<'txn, S: KVStore + 'txn, C: Catalog + ?Sized, T: SqlTxn<'txn, S>>(
57 txn: &mut T,
58 catalog: &C,
59 value: &SqlValue,
60 subquery: &LogicalPlan,
61 negated: bool,
62) -> Result<SqlValue> {
63 execute_in_subquery_with_outer(txn, catalog, value, subquery, negated, None)
64}
65
66pub(crate) fn execute_in_subquery_with_outer<
67 'txn,
68 S: KVStore + 'txn,
69 C: Catalog + ?Sized,
70 T: SqlTxn<'txn, S>,
71>(
72 txn: &mut T,
73 catalog: &C,
74 value: &SqlValue,
75 subquery: &LogicalPlan,
76 negated: bool,
77 outer: Option<&Row>,
78) -> Result<SqlValue> {
79 let rows = execute_subquery_rows_with_outer(txn, catalog, subquery, outer)?;
80 let mut unknown = false;
81 let matched = semi_join_probe(&rows, |row| {
82 let Some(candidate) = row.first() else {
83 return Ok(false);
84 };
85 if matches!(candidate, SqlValue::Null) || matches!(value, SqlValue::Null) {
86 unknown = true;
87 return Ok(false);
88 }
89 compare_values(
90 crate::ast::expr::BinaryOp::Eq,
91 value.clone(),
92 candidate.clone(),
93 )
94 })?;
95 if matched {
96 return Ok(SqlValue::Boolean(!negated));
97 }
98 if unknown {
99 return Ok(SqlValue::Null);
101 }
102 Ok(SqlValue::Boolean(negated))
103}
104
105pub fn execute_exists<'txn, S: KVStore + 'txn, C: Catalog + ?Sized, T: SqlTxn<'txn, S>>(
107 txn: &mut T,
108 catalog: &C,
109 subquery: &LogicalPlan,
110) -> Result<bool> {
111 execute_exists_with_outer(txn, catalog, subquery, false, None)
112}
113
114pub(crate) fn execute_exists_with_outer<
115 'txn,
116 S: KVStore + 'txn,
117 C: Catalog + ?Sized,
118 T: SqlTxn<'txn, S>,
119>(
120 txn: &mut T,
121 catalog: &C,
122 subquery: &LogicalPlan,
123 negated: bool,
124 outer: Option<&Row>,
125) -> Result<bool> {
126 let rows = execute_subquery_rows_with_outer(txn, catalog, subquery, outer)?;
127 let exists = semi_join_probe(&rows, |_| Ok::<bool, ExecutorError>(true))?;
128 Ok(if negated { !exists } else { exists })
129}
130
131pub(crate) fn evaluate_expr_with_subqueries<
132 'txn,
133 S: KVStore + 'txn,
134 C: Catalog + ?Sized,
135 T: SqlTxn<'txn, S>,
136>(
137 txn: &mut T,
138 catalog: &C,
139 expr: &TypedExpr,
140 row: &Row,
141) -> Result<SqlValue> {
142 match &expr.kind {
143 TypedExprKind::ScalarSubquery(subquery) => {
144 execute_scalar_subquery_with_outer(txn, catalog, subquery, Some(row))
145 }
146 TypedExprKind::InSubquery {
147 expr,
148 subquery,
149 negated,
150 } => {
151 let value = evaluate_expr_with_subqueries(txn, catalog, expr, row)?;
152 execute_in_subquery_with_outer(txn, catalog, &value, subquery, *negated, Some(row))
153 }
154 TypedExprKind::Exists { subquery, negated } => {
155 execute_exists_with_outer(txn, catalog, subquery, *negated, Some(row))
156 .map(SqlValue::Boolean)
157 }
158 TypedExprKind::Quantified {
159 expr,
160 op,
161 quantifier,
162 subquery,
163 } => {
164 let value = evaluate_expr_with_subqueries(txn, catalog, expr, row)?;
165 execute_quantified_with_outer(
166 txn,
167 catalog,
168 value,
169 *op,
170 *quantifier,
171 subquery,
172 Some(row),
173 )
174 .map(SqlValue::Boolean)
175 }
176 TypedExprKind::BinaryOp { left, op, right } if contains_subquery(expr) => {
177 let left = evaluate_expr_with_subqueries(txn, catalog, left, row)?;
178 let right = evaluate_expr_with_subqueries(txn, catalog, right, row)?;
179 crate::executor::evaluator::binary_op::eval_binary_values(op, left, right)
180 }
181 _ => {
182 let ctx = EvalContext::new(&row.values);
183 crate::executor::evaluator::evaluate(expr, &ctx)
184 }
185 }
186}
187
188pub(crate) fn contains_subquery(expr: &TypedExpr) -> bool {
189 match &expr.kind {
190 TypedExprKind::ScalarSubquery(_)
191 | TypedExprKind::InSubquery { .. }
192 | TypedExprKind::Exists { .. }
193 | TypedExprKind::Quantified { .. } => true,
194 TypedExprKind::BinaryOp { left, right, .. } => {
195 contains_subquery(left) || contains_subquery(right)
196 }
197 TypedExprKind::UnaryOp { operand, .. } => contains_subquery(operand),
198 TypedExprKind::FunctionCall { args, .. } => args.iter().any(contains_subquery),
199 TypedExprKind::Cast { expr, .. } | TypedExprKind::IsNull { expr, .. } => {
200 contains_subquery(expr)
201 }
202 TypedExprKind::Between {
203 expr, low, high, ..
204 } => contains_subquery(expr) || contains_subquery(low) || contains_subquery(high),
205 TypedExprKind::Like {
206 expr,
207 pattern,
208 escape,
209 ..
210 } => {
211 contains_subquery(expr)
212 || contains_subquery(pattern)
213 || escape.as_deref().is_some_and(contains_subquery)
214 }
215 TypedExprKind::InList { expr, list, .. } => {
216 contains_subquery(expr) || list.iter().any(contains_subquery)
217 }
218 TypedExprKind::Literal(_)
219 | TypedExprKind::ColumnRef { .. }
220 | TypedExprKind::VectorLiteral(_) => false,
221 }
222}
223
224pub(crate) fn plan_contains_subquery(plan: &LogicalPlan) -> bool {
231 fn projection_contains_subquery(projection: &crate::planner::typed_expr::Projection) -> bool {
232 match projection {
233 crate::planner::typed_expr::Projection::All(_) => false,
234 crate::planner::typed_expr::Projection::Columns(cols) => {
235 cols.iter().any(|col| contains_subquery(&col.expr))
236 }
237 }
238 }
239
240 match plan {
241 LogicalPlan::Scan { projection, .. } => projection_contains_subquery(projection),
242 LogicalPlan::Filter { input, predicate } => {
243 contains_subquery(predicate) || plan_contains_subquery(input)
244 }
245 LogicalPlan::Project { input, projection } => {
246 projection_contains_subquery(projection) || plan_contains_subquery(input)
247 }
248 LogicalPlan::Join {
249 left,
250 right,
251 condition,
252 ..
253 } => {
254 condition.as_ref().is_some_and(contains_subquery)
255 || plan_contains_subquery(left)
256 || plan_contains_subquery(right)
257 }
258 LogicalPlan::Aggregate {
259 input,
260 group_keys,
261 aggregates,
262 having,
263 projection,
264 } => {
265 group_keys.iter().any(contains_subquery)
266 || aggregates
267 .iter()
268 .any(|agg| agg.arg.as_ref().is_some_and(contains_subquery))
269 || having.as_ref().is_some_and(contains_subquery)
270 || projection_contains_subquery(projection)
271 || plan_contains_subquery(input)
272 }
273 LogicalPlan::Sort { input, order_by } => {
274 order_by.iter().any(|sort| contains_subquery(&sort.expr))
275 || plan_contains_subquery(input)
276 }
277 LogicalPlan::Limit { input, .. } => plan_contains_subquery(input),
278 _ => false,
280 }
281}
282
283fn execute_quantified_with_outer<
284 'txn,
285 S: KVStore + 'txn,
286 C: Catalog + ?Sized,
287 T: SqlTxn<'txn, S>,
288>(
289 txn: &mut T,
290 catalog: &C,
291 value: SqlValue,
292 op: crate::ast::expr::BinaryOp,
293 quantifier: Quantifier,
294 subquery: &LogicalPlan,
295 outer: Option<&Row>,
296) -> Result<bool> {
297 let rows = execute_subquery_rows_with_outer(txn, catalog, subquery, outer)?;
298 if rows.is_empty() {
299 return Ok(matches!(quantifier, Quantifier::All));
300 }
301 Ok(match quantifier {
302 Quantifier::Any => semi_join_probe(&rows, |row| {
303 let Some(candidate) = row.first() else {
304 return Ok::<bool, ExecutorError>(false);
305 };
306 compare_values(op, value.clone(), candidate.clone())
307 })?,
308 Quantifier::All => {
309 let has_non_match = semi_join_probe(&rows, |row| {
310 let Some(candidate) = row.first() else {
311 return Ok::<bool, ExecutorError>(false);
312 };
313 Ok(!compare_values(op, value.clone(), candidate.clone())?)
314 })?;
315 !has_non_match
316 }
317 })
318}
319
320fn execute_subquery_rows_with_outer<
321 'txn,
322 S: KVStore + 'txn,
323 C: Catalog + ?Sized,
324 T: SqlTxn<'txn, S>,
325>(
326 txn: &mut T,
327 catalog: &C,
328 subquery: &LogicalPlan,
329 outer: Option<&Row>,
330) -> Result<Vec<Vec<SqlValue>>> {
331 if outer.is_none() {
332 let mut cache = materialize_cache();
333 return cache.get_or_try_insert_with((), || {
334 nested_scan(|| {
335 super::execute_query_result_with_outer(txn, catalog, subquery.clone(), outer)
336 .map(|result| result.rows)
337 })
338 });
339 }
340
341 nested_scan(|| {
342 super::execute_query_result_with_outer(txn, catalog, subquery.clone(), outer)
343 .map(|result| result.rows)
344 })
345}
346
347fn compare_values(op: crate::ast::expr::BinaryOp, left: SqlValue, right: SqlValue) -> Result<bool> {
348 match crate::executor::evaluator::binary_op::eval_binary_values(&op, left, right)? {
349 SqlValue::Boolean(value) => Ok(value),
350 other => Err(ExecutorError::Evaluation(EvaluationError::TypeMismatch {
351 expected: "Boolean".into(),
352 actual: other.type_name().into(),
353 })),
354 }
355}
356
357#[cfg(test)]
358mod tests {
359 use super::plan_contains_subquery;
360 use crate::catalog::{Catalog, ColumnMetadata, MemoryCatalog, TableMetadata};
361 use crate::dialect::AlopexDialect;
362 use crate::parser::Parser;
363 use crate::planner::Planner;
364 use crate::planner::logical_plan::LogicalPlan;
365 use crate::planner::types::ResolvedType;
366
367 fn plan_select(sql: &str) -> LogicalPlan {
368 let mut catalog = MemoryCatalog::new();
369 catalog
370 .create_table(TableMetadata::new(
371 "users",
372 vec![
373 ColumnMetadata::new("id", ResolvedType::Integer),
374 ColumnMetadata::new("name", ResolvedType::Text),
375 ],
376 ))
377 .unwrap();
378 catalog
379 .create_table(TableMetadata::new(
380 "orders",
381 vec![
382 ColumnMetadata::new("id", ResolvedType::Integer),
383 ColumnMetadata::new("user_id", ResolvedType::Integer),
384 ColumnMetadata::new("total", ResolvedType::Integer),
385 ],
386 ))
387 .unwrap();
388 let statements = Parser::parse_sql(&AlopexDialect, sql).expect("parse sql");
389 assert_eq!(statements.len(), 1, "expected single statement");
390 Planner::new(&catalog)
391 .plan(&statements[0])
392 .expect("plan sql")
393 }
394
395 #[test]
396 fn detects_subquery_in_join_on_condition() {
397 let plan = plan_select(
398 "SELECT users.name FROM users JOIN orders ON users.id = orders.user_id AND orders.user_id IN (SELECT orders.user_id FROM orders)",
399 );
400 assert!(plan_contains_subquery(&plan));
401 }
402
403 #[test]
404 fn detects_subquery_in_having() {
405 let plan = plan_select(
406 "SELECT orders.user_id, COUNT(*) FROM orders GROUP BY orders.user_id HAVING COUNT(*) > (SELECT MIN(orders.total) FROM orders)",
407 );
408 assert!(plan_contains_subquery(&plan));
409 }
410
411 #[test]
412 fn detects_subquery_in_order_by() {
413 let plan = plan_select(
414 "SELECT users.name FROM users ORDER BY (SELECT COUNT(*) FROM orders WHERE orders.user_id = users.id)",
415 );
416 assert!(plan_contains_subquery(&plan));
417 }
418
419 #[test]
420 fn plan_without_subquery_is_not_detected() {
421 let plan =
422 plan_select("SELECT users.name FROM users WHERE users.id = 1 ORDER BY users.name");
423 assert!(!plan_contains_subquery(&plan));
424 }
425}