Skip to main content

alopex_sql/executor/query/
subquery.rs

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
11/// Execute scalar subquery.
12pub 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
50/// Execute IN subquery.
51///
52/// Implements SQL three-valued logic: when no row matches but the subquery
53/// result contains NULL (or the probe value itself is NULL and the result is
54/// non-empty), the comparison is UNKNOWN and `SqlValue::Null` is returned.
55/// Consequently `NOT IN` over a NULL-containing result never yields TRUE.
56pub 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        // UNKNOWN stays UNKNOWN under NOT.
100        return Ok(SqlValue::Null);
101    }
102    Ok(SqlValue::Boolean(negated))
103}
104
105/// Execute EXISTS subquery.
106pub 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
224/// Returns true if any expression in the query plan contains a subquery.
225///
226/// The streaming pipeline cannot evaluate subqueries because subquery
227/// execution needs transaction access, which the streaming iterators borrow
228/// exclusively. Plans containing subqueries must therefore be routed to the
229/// materializing execution path (see `build_streaming_pipeline_with_policy`).
230pub(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        // DML/DDL plans are never executed through the query pipelines.
279        _ => 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}