Skip to main content

nodedb_sql/planner/lateral/
plan.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! LATERAL join planning: classify correlation shape and emit the appropriate
4//! `SqlPlan::LateralTopK` or `SqlPlan::LateralLoop` variant.
5
6use sqlparser::ast;
7
8use super::correlation::analyse_lateral_where;
9use crate::coerce::expr_as_usize_literal;
10use crate::error::{Result, SqlError};
11use crate::parser::normalize::normalize_ident;
12use crate::resolver::expr::convert_expr;
13use crate::temporal::TemporalScope;
14use crate::types::*;
15
16/// The maximum outer-row count allowed for `LateralLoop` queries.
17pub const LATERAL_LOOP_CAP: usize = 100_000;
18
19/// Parameters for [`plan_lateral_join`].
20pub struct LateralJoinArgs<'a> {
21    /// Plan for the driving (outer) side.
22    pub outer_plan: SqlPlan,
23    /// Alias or name of the outer table for correlation detection.
24    pub outer_alias: Option<String>,
25    /// The LATERAL inner subquery.
26    pub subquery: &'a ast::Query,
27    /// Alias given to the LATERAL in the SQL (e.g. `x` in `LATERAL (...) x`).
28    pub lateral_alias: &'a str,
29    /// True when the enclosing join is LEFT JOIN LATERAL (outer rows
30    /// preserved when inner produces no rows).
31    pub left_join: bool,
32    /// SELECT list projection to apply after the lateral.
33    pub outer_projection: Vec<Projection>,
34    pub catalog: &'a dyn SqlCatalog,
35    pub temporal: TemporalScope,
36}
37
38/// Plan a LATERAL subquery join.
39///
40/// Called when the right side of a JOIN (or a comma-separated FROM item) is a
41/// `TableFactor::Derived { lateral: true, .. }`.
42pub fn plan_lateral_join(args: LateralJoinArgs<'_>) -> Result<SqlPlan> {
43    let LateralJoinArgs {
44        outer_plan,
45        outer_alias,
46        subquery,
47        lateral_alias,
48        left_join,
49        outer_projection,
50        catalog,
51        temporal,
52    } = args;
53    let select = match subquery.body.as_ref() {
54        sqlparser::ast::SetExpr::Select(s) => s,
55        _ => {
56            return Err(SqlError::Unsupported {
57                detail: "LATERAL subquery body must be a SELECT".into(),
58            });
59        }
60    };
61
62    let outer_alias_str = outer_alias.as_deref().unwrap_or("").to_string();
63
64    let analysis = analyse_lateral_where(subquery, &outer_alias_str);
65
66    // Determine if this is the equi-correlated + TopK shape:
67    //   - At least one equi-key correlation.
68    //   - A LIMIT k on the subquery.
69    //   - No non-equi correlations (those require LateralLoop).
70    let has_equi = !analysis.equi_keys.is_empty();
71    let inner_limit = limit_from_query(subquery);
72    let is_top_k = has_equi && inner_limit.is_some() && analysis.non_equi.is_empty();
73
74    if is_top_k {
75        plan_lateral_top_k(LateralTopKPlanArgs {
76            outer_plan,
77            outer_alias,
78            select,
79            subquery,
80            equi_keys: analysis.equi_keys,
81            inner_limit: inner_limit.expect("checked above"),
82            lateral_alias,
83            left_join,
84            outer_projection,
85        })
86    } else if has_equi && analysis.non_equi.is_empty() {
87        // Equi-correlated, no LIMIT: rewrite as a regular hash join.
88        //
89        // The equi correlations become the join keys. The inner side is a bare
90        // scan of the inner collection carrying only the residual (equi-stripped)
91        // WHERE; the outer alias never leaks into inner name resolution. The join
92        // executor scans the inner collection by name, so the join key columns
93        // are available on the merged rows.
94        let inner_plan = build_inner_scan(select, analysis.remaining, catalog, temporal)?;
95        let equi_on: Vec<(String, String)> = analysis
96            .equi_keys
97            .into_iter()
98            .map(|c| (c.outer_col, c.inner_col))
99            .collect();
100        Ok(SqlPlan::Join {
101            left: Box::new(outer_plan),
102            right: Box::new(inner_plan),
103            on: equi_on,
104            join_type: if left_join {
105                JoinType::Left
106            } else {
107                JoinType::Inner
108            },
109            condition: None,
110            limit: None,
111            projection: outer_projection,
112            filters: Vec::new(),
113        })
114    } else {
115        // General correlation — LateralLoop.
116        //
117        // The inner side is a bare scan of the inner collection carrying the
118        // residual WHERE (`analysis.remaining`): non-correlated predicates plus
119        // any non-equi correlated predicates (e.g. `e.log_time > u.created_at`).
120        // The latter lower to `*Column` scan filters (`GtColumn`, ...) whose
121        // outer operand is bound per outer row by the Data Plane executor via
122        // `bind_outer_values`. Equi correlations are applied separately through
123        // `correlation_predicates`, so they are excluded from `remaining` and
124        // not duplicated here. The subquery is not routed through `plan_query`
125        // because its WHERE references the outer alias, which is not resolvable
126        // in the inner FROM scope.
127        let inner_plan = build_inner_scan(select, analysis.remaining, catalog, temporal)?;
128        let correlation_predicates: Vec<(String, String)> = analysis
129            .equi_keys
130            .iter()
131            .map(|c| (c.inner_col.clone(), c.outer_col.clone()))
132            .collect();
133        Ok(SqlPlan::LateralLoop {
134            outer: Box::new(outer_plan),
135            outer_alias,
136            inner: Box::new(inner_plan),
137            correlation_predicates,
138            lateral_alias: lateral_alias.to_string(),
139            projection: outer_projection,
140            outer_row_cap: LATERAL_LOOP_CAP,
141            left_join,
142        })
143    }
144}
145
146/// Build the inner `SqlPlan::Scan` for a LATERAL join.
147///
148/// The scan carries the residual WHERE as filters. Correlated non-equi
149/// predicates survive as column-vs-column comparisons and lower to
150/// runtime-bound `*Column` filters downstream; the executor binds the outer
151/// operand per outer row.
152fn build_inner_scan(
153    select: &sqlparser::ast::Select,
154    residual_where: Option<ast::Expr>,
155    catalog: &dyn SqlCatalog,
156    temporal: TemporalScope,
157) -> Result<SqlPlan> {
158    let inner_collection = extract_inner_collection(select)?;
159    let inner_alias = extract_inner_alias(select);
160    let inner_info = catalog
161        .get_collection(nodedb_types::DatabaseId::DEFAULT, &inner_collection)?
162        .ok_or_else(|| SqlError::UnknownTable {
163            name: inner_collection.clone(),
164        })?;
165    let filters = match &residual_where {
166        Some(expr) => crate::planner::select::convert_where_to_filters(expr)?,
167        None => Vec::new(),
168    };
169    Ok(SqlPlan::Scan {
170        collection: inner_collection,
171        alias: inner_alias,
172        engine: inner_info.engine,
173        filters,
174        projection: Vec::new(),
175        sort_keys: Vec::new(),
176        limit: None,
177        offset: 0,
178        distinct: false,
179        window_functions: Vec::new(),
180        temporal,
181    })
182}
183
184/// Extract the alias of the single-table inner SELECT, if present.
185fn extract_inner_alias(select: &sqlparser::ast::Select) -> Option<String> {
186    let from = select.from.first()?;
187    match &from.relation {
188        ast::TableFactor::Table { alias, .. } => alias.as_ref().map(|a| normalize_ident(&a.name)),
189        _ => None,
190    }
191}
192
193/// Parameters for [`plan_lateral_top_k`].
194struct LateralTopKPlanArgs<'a> {
195    outer_plan: SqlPlan,
196    outer_alias: Option<String>,
197    select: &'a sqlparser::ast::Select,
198    subquery: &'a ast::Query,
199    equi_keys: Vec<super::correlation::CorrelationEq>,
200    inner_limit: usize,
201    lateral_alias: &'a str,
202    left_join: bool,
203    outer_projection: Vec<Projection>,
204}
205
206/// Plan the `LateralTopK` variant: equi-correlated + ORDER BY + LIMIT k.
207fn plan_lateral_top_k(args: LateralTopKPlanArgs<'_>) -> Result<SqlPlan> {
208    let LateralTopKPlanArgs {
209        outer_plan,
210        outer_alias,
211        select,
212        subquery,
213        equi_keys,
214        inner_limit,
215        lateral_alias,
216        left_join,
217        outer_projection,
218    } = args;
219    // Build a bare inner Scan without correlation filters (those are injected
220    // at runtime per outer row).
221    let inner_collection = extract_inner_collection(select)?;
222    let inner_filters = inner_non_correlated_filters(select, outer_alias.as_deref().unwrap_or(""))?;
223
224    // Extract ORDER BY from the inner subquery.
225    // For LATERAL inner scans we only need simple column-expression sort keys;
226    // the full search-trigger machinery (vector/hybrid search) is not applicable
227    // here, so we convert expressions directly.
228    let inner_order_by = if let Some(order_by) = &subquery.order_by {
229        match &order_by.kind {
230            ast::OrderByKind::Expressions(exprs) => exprs
231                .iter()
232                .filter_map(|o| {
233                    convert_expr(&o.expr).ok().map(|expr| SortKey {
234                        expr,
235                        ascending: o.options.asc.unwrap_or(true),
236                        nulls_first: o.options.nulls_first.unwrap_or(false),
237                    })
238                })
239                .collect(),
240            ast::OrderByKind::All(_) => Vec::new(),
241        }
242    } else {
243        Vec::new()
244    };
245
246    let correlation_keys: Vec<(String, String)> = equi_keys
247        .into_iter()
248        .map(|c| (c.outer_col, c.inner_col))
249        .collect();
250
251    Ok(SqlPlan::LateralTopK {
252        outer: Box::new(outer_plan),
253        outer_alias,
254        inner_collection,
255        inner_filters,
256        inner_order_by,
257        inner_limit,
258        correlation_keys,
259        lateral_alias: lateral_alias.to_string(),
260        projection: outer_projection,
261        left_join,
262    })
263}
264
265/// Extract the collection name from a single-table inner SELECT.
266fn extract_inner_collection(select: &sqlparser::ast::Select) -> Result<String> {
267    let from = select.from.first().ok_or_else(|| SqlError::Unsupported {
268        detail: "LATERAL subquery must have a FROM clause".into(),
269    })?;
270    match &from.relation {
271        ast::TableFactor::Table { name, .. } => {
272            crate::parser::normalize::normalize_object_name_checked(name)
273        }
274        _ => Err(SqlError::Unsupported {
275            detail: "LATERAL LateralTopK subquery must reference a plain table".into(),
276        }),
277    }
278}
279
280/// Extract filters from the inner SELECT that do NOT reference the outer alias.
281fn inner_non_correlated_filters(
282    select: &sqlparser::ast::Select,
283    outer_alias: &str,
284) -> Result<Vec<Filter>> {
285    let Some(where_expr) = &select.selection else {
286        return Ok(Vec::new());
287    };
288    let remaining = strip_outer_refs(where_expr, outer_alias);
289    match remaining {
290        Some(expr) => crate::planner::select::convert_where_to_filters(&expr),
291        None => Ok(Vec::new()),
292    }
293}
294
295/// Remove all predicates referencing `outer_alias` from a WHERE expression.
296fn strip_outer_refs(expr: &ast::Expr, outer_alias: &str) -> Option<ast::Expr> {
297    match expr {
298        ast::Expr::BinaryOp {
299            left,
300            op: ast::BinaryOperator::And,
301            right,
302        } => {
303            let l = strip_outer_refs(left, outer_alias);
304            let r = strip_outer_refs(right, outer_alias);
305            match (l, r) {
306                (None, None) => None,
307                (Some(e), None) | (None, Some(e)) => Some(e),
308                (Some(l), Some(r)) => Some(ast::Expr::BinaryOp {
309                    left: Box::new(l),
310                    op: ast::BinaryOperator::And,
311                    right: Box::new(r),
312                }),
313            }
314        }
315        ast::Expr::BinaryOp { left, right, .. } => {
316            if refs_outer(left, outer_alias) || refs_outer(right, outer_alias) {
317                None
318            } else {
319                Some(expr.clone())
320            }
321        }
322        ast::Expr::Nested(inner) => strip_outer_refs(inner, outer_alias),
323        _ => Some(expr.clone()),
324    }
325}
326
327fn refs_outer(expr: &ast::Expr, outer_alias: &str) -> bool {
328    match expr {
329        ast::Expr::CompoundIdentifier(parts) if parts.len() == 2 => {
330            normalize_ident(&parts[0]).eq_ignore_ascii_case(outer_alias)
331        }
332        ast::Expr::BinaryOp { left, right, .. } => {
333            refs_outer(left, outer_alias) || refs_outer(right, outer_alias)
334        }
335        _ => false,
336    }
337}
338
339/// Extract the LIMIT value from a query.
340fn limit_from_query(query: &ast::Query) -> Option<usize> {
341    match &query.limit_clause {
342        Some(ast::LimitClause::LimitOffset { limit, .. }) => {
343            limit.as_ref().and_then(expr_as_usize_literal)
344        }
345        Some(ast::LimitClause::OffsetCommaLimit { limit, .. }) => {
346            Some(expr_as_usize_literal(limit).unwrap_or(0))
347        }
348        None => None,
349    }
350}
351
352/// Extract a LATERAL alias from a `TableFactor::Derived`.
353pub fn lateral_alias_from_factor(factor: &ast::TableFactor) -> Option<String> {
354    match factor {
355        ast::TableFactor::Derived { alias, .. } => alias.as_ref().map(|a| normalize_ident(&a.name)),
356        _ => None,
357    }
358}
359
360/// True when a `TableFactor` is a LATERAL derived subquery.
361pub fn is_lateral_derived(factor: &ast::TableFactor) -> bool {
362    matches!(factor, ast::TableFactor::Derived { lateral: true, .. })
363}
364
365/// Extract the subquery from a `TableFactor::Derived`.
366pub fn subquery_from_factor(factor: &ast::TableFactor) -> Option<&ast::Query> {
367    match factor {
368        ast::TableFactor::Derived { subquery, .. } => Some(subquery),
369        _ => None,
370    }
371}