uqa_execution/
join_output.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Positional output shaping for `JOIN ... USING` and `NATURAL JOIN`.
8
9use uqa_core::Value;
10use uqa_sql::ast::ColumnType;
11use uqa_sql::expr::cast_value;
12
13use crate::{Batch, ColumnIdentity, ExecError, ExecResult, PhysicalOperator, RowSchema};
14
15/// Source of one visible or hidden join-output identity.
16#[derive(Debug, Clone, PartialEq, Eq)]
17pub enum JoinOutputSource {
18    /// Reuse an existing logical input position without copying its value.
19    Input(usize),
20    /// Apply an implicit binder-selected coercion to one input position.
21    Cast { input: usize, ty: ColumnType },
22    /// SQL `COALESCE(left::type, right::type)` over two logical input
23    /// positions. This is required only for a merged column of `FULL JOIN`.
24    Coalesce {
25        left: usize,
26        right: usize,
27        ty: ColumnType,
28    },
29}
30
31/// Reorder and merge join columns while retaining the joined physical row.
32/// Inner, left, and right joins are schema-only remaps; a full join appends
33/// only the merged values that cannot be represented by one existing slot.
34pub struct JoinOutput<'a> {
35    child: Box<dyn PhysicalOperator + 'a>,
36    schema: RowSchema,
37    computed: Vec<JoinOutputSource>,
38}
39
40impl<'a> JoinOutput<'a> {
41    /// Compile the positional schema of a qualified join without constructing or executing an operator. Binders use this to expose the exact same merged-column identities and hidden qualified aliases as execution.
42    pub fn try_schema(
43        input: &RowSchema,
44        columns: &[(String, ColumnIdentity, JoinOutputSource)],
45        aliases: &[(ColumnIdentity, JoinOutputSource)],
46    ) -> ExecResult<RowSchema> {
47        compile_layout(input, columns, aliases).map(|(schema, _)| schema)
48    }
49
50    pub fn try_new(
51        child: Box<dyn PhysicalOperator + 'a>,
52        columns: Vec<(String, ColumnIdentity, JoinOutputSource)>,
53        aliases: Vec<(ColumnIdentity, JoinOutputSource)>,
54    ) -> ExecResult<Self> {
55        let (schema, computed) = compile_layout(child.row_schema(), &columns, &aliases)?;
56        Ok(Self {
57            child,
58            schema,
59            computed,
60        })
61    }
62}
63
64fn compile_layout(
65    input: &RowSchema,
66    columns: &[(String, ColumnIdentity, JoinOutputSource)],
67    aliases: &[(ColumnIdentity, JoinOutputSource)],
68) -> ExecResult<(RowSchema, Vec<JoinOutputSource>)> {
69    let input_width = input.len();
70    let mut computed = Vec::<JoinOutputSource>::new();
71    for source in columns
72        .iter()
73        .map(|(_, _, source)| source)
74        .chain(aliases.iter().map(|(_, source)| source))
75    {
76        match source {
77            JoinOutputSource::Input(position) if *position >= input_width => {
78                return Err(ExecError::Other(format!(
79                    "join output input position {position} is outside width {input_width}"
80                )));
81            }
82            JoinOutputSource::Cast { input, .. } if *input >= input_width => {
83                return Err(ExecError::Other(format!(
84                    "join output cast position {input} is outside width {input_width}"
85                )));
86            }
87            JoinOutputSource::Coalesce { left, right, .. }
88                if *left >= input_width || *right >= input_width =>
89            {
90                return Err(ExecError::Other(format!(
91                    "join output coalesce positions ({left}, {right}) are outside width {input_width}"
92                )));
93            }
94            source @ (JoinOutputSource::Cast { .. } | JoinOutputSource::Coalesce { .. }) => {
95                if !computed.contains(source) {
96                    computed.push(source.clone());
97                }
98            }
99            JoinOutputSource::Input(_) => {}
100        }
101    }
102
103    let computed_names = (0..computed.len())
104        .map(|index| format!("\0uqa.join_using.{index}"))
105        .collect::<Vec<_>>();
106    let computed_columns = computed_names
107        .iter()
108        .cloned()
109        .zip(computed.iter().map(source_type))
110        .collect::<Vec<_>>();
111    let intermediate = RowSchema::append_typed(input, &computed_columns);
112    let source_position = |source: &JoinOutputSource| -> usize {
113        match source {
114            JoinOutputSource::Input(position) => *position,
115            JoinOutputSource::Cast { .. } | JoinOutputSource::Coalesce { .. } => {
116                let index = computed
117                    .iter()
118                    .position(|candidate| candidate == source)
119                    .expect("computed join output source was registered");
120                intermediate
121                    .position(&computed_names[index])
122                    .expect("computed join output column exists")
123            }
124        }
125    };
126    let columns = columns
127        .iter()
128        .map(|(name, identity, source)| {
129            let ty = match source {
130                JoinOutputSource::Input(position) => intermediate.column_type(*position).cloned(),
131                JoinOutputSource::Cast { .. } | JoinOutputSource::Coalesce { .. } => {
132                    source_type(source)
133                }
134            };
135            (name.clone(), identity.clone(), source_position(source), ty)
136        })
137        .collect::<Vec<_>>();
138    let aliases = aliases
139        .iter()
140        .map(|(name, source)| (name.clone(), source_position(source)))
141        .collect::<Vec<_>>();
142    let schema = RowSchema::remap_typed_identities(&intermediate, &columns, &aliases);
143    Ok((schema, computed))
144}
145
146impl PhysicalOperator for JoinOutput<'_> {
147    fn row_schema(&self) -> &RowSchema {
148        &self.schema
149    }
150
151    fn estimated_cardinality(&self) -> Option<u64> {
152        self.child.estimated_cardinality()
153    }
154
155    fn open(&mut self) -> ExecResult<()> {
156        self.child.open()
157    }
158
159    fn next(&mut self) -> ExecResult<Option<Batch>> {
160        let Some(batch) = self.child.next()? else {
161            return Ok(None);
162        };
163        if self.computed.is_empty() {
164            return Ok(Some(Batch::from_physical_rows(
165                self.schema.clone(),
166                batch.rows,
167            )));
168        }
169        let rows = batch
170            .rows
171            .into_iter()
172            .map(|row| {
173                let values = {
174                    let view = batch.schema.view(&row);
175                    self.computed
176                        .iter()
177                        .map(|source| evaluate_source(source, &view))
178                        .collect::<ExecResult<Vec<_>>>()?
179                };
180                Ok(row.append_values(values))
181            })
182            .collect::<ExecResult<Vec<_>>>()?;
183        Ok(Some(Batch::from_physical_rows(self.schema.clone(), rows)))
184    }
185
186    fn close(&mut self) -> ExecResult<()> {
187        self.child.close()
188    }
189}
190
191fn source_type(source: &JoinOutputSource) -> Option<ColumnType> {
192    match source {
193        JoinOutputSource::Input(_) => None,
194        JoinOutputSource::Cast { ty, .. } | JoinOutputSource::Coalesce { ty, .. } => {
195            Some(ty.clone())
196        }
197    }
198}
199
200fn evaluate_source(
201    source: &JoinOutputSource,
202    view: &crate::PhysicalRowView<'_>,
203) -> ExecResult<Value> {
204    match source {
205        JoinOutputSource::Input(position) => {
206            Ok(view.value_at(*position).unwrap_or(&Value::Null).clone())
207        }
208        JoinOutputSource::Cast { input, ty } => cast_value(
209            view.value_at(*input).unwrap_or(&Value::Null),
210            &ty.sql_name(),
211        )
212        .map_err(ExecError::from),
213        JoinOutputSource::Coalesce { left, right, ty } => {
214            let target = ty.sql_name();
215            let left = cast_value(view.value_at(*left).unwrap_or(&Value::Null), &target)?;
216            if !matches!(left, Value::Null) {
217                return Ok(left);
218            }
219            cast_value(view.value_at(*right).unwrap_or(&Value::Null), &target)
220                .map_err(ExecError::from)
221        }
222    }
223}
224
225#[cfg(test)]
226mod tests {
227    use std::collections::BTreeMap;
228
229    use super::*;
230    use crate::physical::run_to_rows;
231    use crate::TableScan;
232    use uqa_sql::expr::RowLookup;
233
234    #[test]
235    fn schema_only_merge_reuses_input_slots_and_keeps_qualified_aliases() {
236        let child = TableScan::from_rows(
237            vec![
238                "l.id".into(),
239                "l.name".into(),
240                "r.id".into(),
241                "r.note".into(),
242            ],
243            vec![BTreeMap::from([
244                ("l.id".into(), Value::Int(1)),
245                ("l.name".into(), Value::Str("left".into())),
246                ("r.id".into(), Value::Int(1)),
247                ("r.note".into(), Value::Str("right".into())),
248            ])],
249        );
250        let mut output = JoinOutput::try_new(
251            Box::new(child),
252            vec![
253                (
254                    "id".into(),
255                    ColumnIdentity::unqualified("id"),
256                    JoinOutputSource::Input(0),
257                ),
258                (
259                    "name".into(),
260                    ColumnIdentity::qualified("l", "name"),
261                    JoinOutputSource::Input(1),
262                ),
263                (
264                    "note".into(),
265                    ColumnIdentity::qualified("r", "note"),
266                    JoinOutputSource::Input(3),
267                ),
268            ],
269            vec![
270                (
271                    ColumnIdentity::qualified("l", "id"),
272                    JoinOutputSource::Input(0),
273                ),
274                (
275                    ColumnIdentity::qualified("r", "id"),
276                    JoinOutputSource::Input(2),
277                ),
278            ],
279        )
280        .unwrap();
281        output.open().unwrap();
282        let batch = output.next().unwrap().unwrap();
283        assert_eq!(batch.schema.columns(), ["id", "name", "note"]);
284        let view = batch.schema.view(&batch.rows[0]);
285        assert_eq!(view.qualified_column("l", "id"), Some(&Value::Int(1)));
286        assert_eq!(view.qualified_column("r", "id"), Some(&Value::Int(1)));
287        output.close().unwrap();
288    }
289
290    #[test]
291    fn full_merge_coalesces_only_the_requested_column() {
292        let child = TableScan::from_rows(
293            vec!["l.id".into(), "r.id".into()],
294            vec![BTreeMap::from([
295                ("l.id".into(), Value::Null),
296                ("r.id".into(), Value::Int(3)),
297            ])],
298        );
299        let mut output = JoinOutput::try_new(
300            Box::new(child),
301            vec![(
302                "id".into(),
303                ColumnIdentity::unqualified("id"),
304                JoinOutputSource::Coalesce {
305                    left: 0,
306                    right: 1,
307                    ty: ColumnType::Integer,
308                },
309            )],
310            Vec::new(),
311        )
312        .unwrap();
313        let (_, rows) = run_to_rows(&mut output).unwrap();
314        assert_eq!(rows[0]["id"], Value::Int(3));
315    }
316}
uqa_execution/join_output.rs

uqa_execution/
join_output.rs
