Skip to main content

nodedb_lite/query/query_ops/
lateral_loop.rs

1// SPDX-License-Identifier: Apache-2.0
2//! LateralLoop — nested-loop lateral join with correlated predicates.
3//!
4//! For each outer row, re-scans the inner collection with equality filters
5//! built from `correlation_predicates` (inner_field = outer_field_value).
6//! No sort, no limit per outer row.  Enforces `outer_row_cap`.
7
8use std::collections::HashMap;
9
10use nodedb_physical::physical_plan::PhysicalPlan;
11use nodedb_physical::physical_plan::query::JoinProjection;
12use nodedb_query::scan_filter::FilterOp;
13use nodedb_query::scan_filter::ScanFilter;
14use nodedb_types::result::QueryResult;
15use nodedb_types::value::Value;
16
17use crate::error::LiteError;
18use crate::query::engine::LiteQueryEngine;
19use crate::query::query_ops::joins::common::{
20    apply_filters, apply_projection, decode_filters, maps_to_result, scan_collection,
21};
22use nodedb_sql::types::SqlPlan;
23
24use crate::query::query_ops::lateral_top_k::{execute_nested_plan, prefix_row};
25use crate::storage::engine::StorageEngine;
26
27/// Inputs of a LATERAL nested-loop join over a SQL inner plan.
28pub(crate) struct LateralLoopSqlSpec<'a> {
29    pub outer_alias: &'a str,
30    pub inner_sql: &'a SqlPlan,
31    pub correlation_predicates: &'a [(String, String)],
32    pub lateral_alias: &'a str,
33    pub projection: &'a [JoinProjection],
34    pub left_join: bool,
35    pub outer_row_cap: usize,
36}
37
38/// Inputs of a LATERAL nested-loop join over an inner collection.
39///
40/// `inner_filters` holds msgpack-encoded `MetadataFilter` bytes. An empty
41/// slice means no filter.
42pub struct LateralLoopSpec<'a> {
43    pub outer_alias: &'a str,
44    pub inner_collection: &'a str,
45    pub inner_filters: &'a [u8],
46    pub correlation_predicates: &'a [(String, String)],
47    pub lateral_alias: &'a str,
48    pub projection: &'a [JoinProjection],
49    pub left_join: bool,
50    pub outer_row_cap: usize,
51}
52
53/// SQL-plan-aware entry for `LiteVisitor::lateral_loop`.
54///
55/// Executes `outer_sql` via `engine.execute_plan`, then for each outer row
56/// re-executes `inner_sql` with correlated equality filters injected.
57pub(crate) async fn execute_lateral_loop_sql<S: StorageEngine>(
58    engine: &LiteQueryEngine<S>,
59    permit: Option<&crate::engine::fts::coordinator::TextMutationPermit>,
60    outer_sql: &SqlPlan,
61    spec: LateralLoopSqlSpec<'_>,
62) -> Result<QueryResult, LiteError> {
63    let LateralLoopSqlSpec {
64        outer_alias,
65        inner_sql,
66        correlation_predicates,
67        lateral_alias,
68        projection,
69        left_join,
70        outer_row_cap,
71    } = spec;
72    use crate::query::query_ops::joins::common::rows_to_maps;
73
74    let outer_result = engine.execute_plan_admitted(outer_sql, permit).await?;
75    let outer_rows = rows_to_maps(outer_result);
76
77    if outer_row_cap > 0 && outer_rows.len() > outer_row_cap {
78        return Err(LiteError::Storage {
79            detail: format!(
80                "lateral loop outer_row_cap={outer_row_cap} exceeded: got {} outer rows",
81                outer_rows.len()
82            ),
83        });
84    }
85
86    let mut output: Vec<HashMap<String, Value>> = Vec::new();
87
88    for outer_row in &outer_rows {
89        let outer_prefixed = prefix_row(outer_row, outer_alias);
90
91        let inner_result = engine.execute_plan_admitted(inner_sql, permit).await?;
92        let inner_all = rows_to_maps(inner_result);
93
94        let mut corr_filters: Vec<ScanFilter> = Vec::new();
95        for (inner_field, outer_field) in correlation_predicates {
96            let val = outer_row.get(outer_field).cloned().unwrap_or(Value::Null);
97            corr_filters.push(ScanFilter {
98                field: inner_field.clone(),
99                op: FilterOp::Eq,
100                value: val,
101                clauses: Vec::new(),
102                expr: None,
103            });
104        }
105        let inner_rows = apply_filters(inner_all, &corr_filters)?;
106
107        if inner_rows.is_empty() {
108            if left_join {
109                let projected = apply_projection(vec![outer_prefixed], projection);
110                output.extend(projected);
111            }
112            continue;
113        }
114        for inner_row in &inner_rows {
115            let inner_prefixed = prefix_row(inner_row, lateral_alias);
116            let mut merged = outer_prefixed.clone();
117            merged.extend(inner_prefixed);
118            let projected = apply_projection(vec![merged], projection);
119            output.extend(projected);
120        }
121    }
122
123    Ok(maps_to_result(output))
124}
125
126/// LATERAL nested-loop with correlated predicates injected per outer row.
127///
128/// Runs `outer_plan` once to materialise outer rows, caps them at
129/// `outer_row_cap` (0 = uncapped), then for each outer row builds equality
130/// filters from `correlation_predicates` and scans `inner_collection`.
131/// Emits every matching inner row merged with the outer row.  Supports LEFT
132/// join semantics.
133pub async fn execute_lateral_loop<S: StorageEngine>(
134    engine: &LiteQueryEngine<S>,
135    outer_plan: &PhysicalPlan,
136    spec: LateralLoopSpec<'_>,
137) -> Result<QueryResult, LiteError> {
138    let LateralLoopSpec {
139        outer_alias,
140        inner_collection,
141        inner_filters,
142        correlation_predicates,
143        lateral_alias,
144        projection,
145        left_join,
146        outer_row_cap,
147    } = spec;
148    let outer_rows = execute_nested_plan(engine, outer_plan).await?;
149
150    if outer_row_cap > 0 && outer_rows.len() > outer_row_cap {
151        return Err(LiteError::Storage {
152            detail: format!(
153                "lateral loop outer_row_cap={outer_row_cap} exceeded: got {} outer rows",
154                outer_rows.len()
155            ),
156        });
157    }
158
159    let base_inner_filters = decode_filters(inner_filters)?;
160
161    let mut output: Vec<HashMap<String, Value>> = Vec::new();
162
163    for outer_row in &outer_rows {
164        // Build correlated equality filters.
165        let mut filters = base_inner_filters.clone();
166        for (inner_field, outer_field) in correlation_predicates {
167            let val = outer_row.get(outer_field).cloned().unwrap_or(Value::Null);
168            filters.push(ScanFilter {
169                field: inner_field.clone(),
170                op: FilterOp::Eq,
171                value: val,
172                clauses: Vec::new(),
173                expr: None,
174            });
175        }
176
177        let inner_all = scan_collection(engine, inner_collection).await?;
178        let inner_rows = apply_filters(inner_all, &filters)?;
179
180        if inner_rows.is_empty() {
181            if left_join {
182                let merged = prefix_row(outer_row, outer_alias);
183                output.push(merged);
184            }
185            continue;
186        }
187
188        for inner_row in &inner_rows {
189            let outer_prefixed = prefix_row(outer_row, outer_alias);
190            let inner_prefixed = prefix_row(inner_row, lateral_alias);
191            let mut merged = outer_prefixed;
192            merged.extend(inner_prefixed);
193            output.push(merged);
194        }
195    }
196
197    let projected = apply_projection(output, projection);
198    Ok(maps_to_result(projected))
199}
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204
205    fn customer(cust_id: i64, region: &str) -> HashMap<String, Value> {
206        let mut m = HashMap::new();
207        m.insert("cust_id".to_string(), Value::Integer(cust_id));
208        m.insert("region".to_string(), Value::String(region.to_string()));
209        m
210    }
211
212    fn order(order_id: i64, cust_id: i64, amount: i64) -> HashMap<String, Value> {
213        let mut m = HashMap::new();
214        m.insert("order_id".to_string(), Value::Integer(order_id));
215        m.insert("cust_id".to_string(), Value::Integer(cust_id));
216        m.insert("amount".to_string(), Value::Integer(amount));
217        m
218    }
219
220    /// Correlated WHERE: each customer row triggers an inner scan that
221    /// matches orders with matching cust_id.
222    #[test]
223    fn test_lateral_loop_correlated_where() {
224        let customers = vec![customer(1, "west"), customer(2, "east")];
225        let orders = vec![
226            order(1, 1, 100),
227            order(2, 1, 200),
228            order(3, 1, 150),
229            order(4, 2, 300),
230        ];
231
232        let mut all_output: Vec<HashMap<String, Value>> = Vec::new();
233        for outer_row in &customers {
234            let cust_val = outer_row.get("cust_id").cloned().unwrap_or(Value::Null);
235            let filters = vec![ScanFilter {
236                field: "cust_id".to_string(),
237                op: FilterOp::Eq,
238                value: cust_val,
239                clauses: Vec::new(),
240                expr: None,
241            }];
242            let inner_rows =
243                apply_filters(orders.clone(), &filters).expect("filter evaluation failed");
244            for inner_row in inner_rows {
245                let mut merged = outer_row.clone();
246                merged.extend(inner_row);
247                all_output.push(merged);
248            }
249        }
250
251        // cust 1 → 3 orders; cust 2 → 1 order.
252        assert_eq!(all_output.len(), 4);
253        // All output rows carry the customer's region.
254        assert!(all_output.iter().all(|r| r.contains_key("region")));
255    }
256
257    /// outer_row_cap enforcement: error when cap is exceeded.
258    #[test]
259    fn test_outer_row_cap_error_message() {
260        let outer_rows: Vec<HashMap<String, Value>> = (0..5)
261            .map(|i| {
262                let mut m = HashMap::new();
263                m.insert("id".to_string(), Value::Integer(i));
264                m
265            })
266            .collect();
267        let cap = 3usize;
268        let exceeded = outer_rows.len() > cap;
269        assert!(exceeded, "cap should be exceeded");
270    }
271}