Skip to main content

uqa_execution/scalar/
subquery.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Pull-based scalar-subquery protocol and result consumers.
8
9use uqa_core::Value;
10use uqa_sql::expr::RowLookup;
11use uqa_sql::{ResultRow, SQLError, SQLParam};
12
13use crate::batch::{OwnedPhysicalRow, PhysicalRow, RowSchema};
14
15use super::SubqueryId;
16
17/// Runtime callback for query children referenced by [`SubqueryId`]. The planner owns the actual query-plan arena; execution only needs this stable slot interface.
18pub trait ScalarSubqueryRunner {
19    fn execute_subquery(
20        &self,
21        subquery: SubqueryId,
22        outer_row: Option<&dyn RowLookup>,
23        params: &[SQLParam],
24    ) -> Result<SubqueryResult, SQLError>;
25
26    fn execute_subquery_physical(
27        &self,
28        subquery: SubqueryId,
29        outer_schema: &RowSchema,
30        outer_row: &PhysicalRow,
31        params: &[SQLParam],
32    ) -> Result<SubqueryResult, SQLError> {
33        let outer = outer_schema.view(outer_row);
34        self.execute_subquery(subquery, Some(&outer), params)
35    }
36
37    fn scalar_subquery_value(
38        &self,
39        subquery: SubqueryId,
40        outer_row: Option<&dyn RowLookup>,
41        params: &[SQLParam],
42    ) -> Result<Value, SQLError> {
43        self.execute_subquery(subquery, outer_row, params)?
44            .into_scalar_value()
45    }
46
47    fn scalar_subquery_value_physical(
48        &self,
49        subquery: SubqueryId,
50        outer_schema: &RowSchema,
51        outer_row: &PhysicalRow,
52        params: &[SQLParam],
53    ) -> Result<Value, SQLError> {
54        self.execute_subquery_physical(subquery, outer_schema, outer_row, params)?
55            .into_scalar_value()
56    }
57
58    fn subquery_exists(
59        &self,
60        subquery: SubqueryId,
61        outer_row: Option<&dyn RowLookup>,
62        params: &[SQLParam],
63    ) -> Result<bool, SQLError> {
64        self.execute_subquery(subquery, outer_row, params)?
65            .into_exists()
66    }
67
68    fn subquery_exists_physical(
69        &self,
70        subquery: SubqueryId,
71        outer_schema: &RowSchema,
72        outer_row: &PhysicalRow,
73        params: &[SQLParam],
74    ) -> Result<bool, SQLError> {
75        self.execute_subquery_physical(subquery, outer_schema, outer_row, params)?
76            .into_exists()
77    }
78
79    fn subquery_contains(
80        &self,
81        subquery: SubqueryId,
82        needle: &Value,
83        outer_row: Option<&dyn RowLookup>,
84        params: &[SQLParam],
85    ) -> Result<Option<bool>, SQLError> {
86        self.execute_subquery(subquery, outer_row, params)?
87            .contains(needle)
88    }
89
90    fn subquery_contains_physical(
91        &self,
92        subquery: SubqueryId,
93        needle: &Value,
94        outer_schema: &RowSchema,
95        outer_row: &PhysicalRow,
96        params: &[SQLParam],
97    ) -> Result<Option<bool>, SQLError> {
98        self.execute_subquery_physical(subquery, outer_schema, outer_row, params)?
99            .contains(needle)
100    }
101}
102
103/// Pull-based scalar-subquery result. Scalar, EXISTS, and IN consumers never need to materialize the complete child relation: they respectively inspect at most two rows, one row, or one row at a time.
104pub struct SubqueryResult {
105    pub columns: Vec<String>,
106    pub rows: Box<dyn Iterator<Item = Result<OwnedPhysicalRow, SQLError>> + Send>,
107}
108
109impl SubqueryResult {
110    pub fn from_rows(columns: Vec<String>, rows: Vec<ResultRow>) -> Self {
111        let schema = RowSchema::new(columns.clone());
112        Self {
113            columns,
114            rows: Box::new(rows.into_iter().map(move |row| {
115                Ok(OwnedPhysicalRow::new(
116                    schema.clone(),
117                    PhysicalRow::from_result_row(&schema, row),
118                ))
119            })),
120        }
121    }
122
123    pub fn into_scalar_value(mut self) -> Result<Value, SQLError> {
124        let Some(first_row) = self.rows.next().transpose()? else {
125            return Ok(Value::Null);
126        };
127        if self.rows.next().transpose()?.is_some() {
128            return Err(SQLError::Routine {
129                sqlstate: "21000".into(),
130                message: "more than one row returned by a subquery used as an expression".into(),
131            });
132        }
133        if self.columns.is_empty() {
134            return Err(SQLError::TypeMismatch(
135                "scalar subquery returned no columns".into(),
136            ));
137        }
138        Ok(first_row
139            .positional_column(0)
140            .cloned()
141            .unwrap_or(Value::Null))
142    }
143
144    pub fn into_exists(mut self) -> Result<bool, SQLError> {
145        Ok(self.rows.next().transpose()?.is_some())
146    }
147
148    pub fn contains(self, needle: &Value) -> Result<Option<bool>, SQLError> {
149        if self.columns.is_empty() {
150            return Ok(Some(false));
151        }
152        let mut saw_row = false;
153        let mut saw_null = false;
154        for row in self.rows {
155            let row = row?;
156            saw_row = true;
157            match row.positional_column(0) {
158                Some(Value::Null) | None => saw_null = true,
159                Some(value) if !matches!(needle, Value::Null) && value == needle => {
160                    return Ok(Some(true));
161                }
162                Some(_) => {}
163            }
164        }
165        Ok(if !saw_row {
166            Some(false)
167        } else if matches!(needle, Value::Null) || saw_null {
168            None
169        } else {
170            Some(false)
171        })
172    }
173}