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_types = computed.iter().map(source_type).collect::<Vec<_>>();
104    let intermediate = RowSchema::append_hidden_typed(input, &computed_types);
105    let source_position = |source: &JoinOutputSource| -> usize {
106        match source {
107            JoinOutputSource::Input(position) => input
108                .physical_slot(*position)
109                .expect("validated join output input position has a physical slot"),
110            JoinOutputSource::Cast { .. } | JoinOutputSource::Coalesce { .. } => {
111                let index = computed
112                    .iter()
113                    .position(|candidate| candidate == source)
114                    .expect("computed join output source was registered");
115                input.physical_width() + index
116            }
117        }
118    };
119    let columns = columns
120        .iter()
121        .map(|(name, identity, source)| {
122            let ty = match source {
123                JoinOutputSource::Input(position) => intermediate.column_type(*position).cloned(),
124                JoinOutputSource::Cast { .. } | JoinOutputSource::Coalesce { .. } => {
125                    source_type(source)
126                }
127            };
128            (name.clone(), identity.clone(), source_position(source), ty)
129        })
130        .collect::<Vec<_>>();
131    let aliases = aliases
132        .iter()
133        .map(|(name, source)| {
134            let ty = match source {
135                JoinOutputSource::Input(position) => input.column_type(*position).cloned(),
136                JoinOutputSource::Cast { .. } | JoinOutputSource::Coalesce { .. } => {
137                    source_type(source)
138                }
139            };
140            (name.clone(), source_position(source), ty)
141        })
142        .collect::<Vec<_>>();
143    let schema = RowSchema::remap_typed_physical_identities(&intermediate, &columns, &aliases);
144    Ok((schema, computed))
145}
146
147impl PhysicalOperator for JoinOutput<'_> {
148    fn row_schema(&self) -> &RowSchema {
149        &self.schema
150    }
151
152    fn estimated_cardinality(&self) -> Option<u64> {
153        self.child.estimated_cardinality()
154    }
155
156    fn open(&mut self) -> ExecResult<()> {
157        self.child.open()
158    }
159
160    fn next(&mut self) -> ExecResult<Option<Batch>> {
161        let Some(batch) = self.child.next()? else {
162            return Ok(None);
163        };
164        if self.computed.is_empty() {
165            return Ok(Some(Batch::from_physical_rows(
166                self.schema.clone(),
167                batch.rows,
168            )));
169        }
170        let rows = batch
171            .rows
172            .into_iter()
173            .map(|row| {
174                let values = {
175                    let view = batch.schema.view(&row);
176                    self.computed
177                        .iter()
178                        .map(|source| evaluate_source(source, &view))
179                        .collect::<ExecResult<Vec<_>>>()?
180                };
181                Ok(row.append_values(values))
182            })
183            .collect::<ExecResult<Vec<_>>>()?;
184        Ok(Some(Batch::from_physical_rows(self.schema.clone(), rows)))
185    }
186
187    fn close(&mut self) -> ExecResult<()> {
188        self.child.close()
189    }
190}
191
192fn source_type(source: &JoinOutputSource) -> Option<ColumnType> {
193    match source {
194        JoinOutputSource::Input(_) => None,
195        JoinOutputSource::Cast { ty, .. } | JoinOutputSource::Coalesce { ty, .. } => {
196            Some(ty.clone())
197        }
198    }
199}
200
201fn evaluate_source(
202    source: &JoinOutputSource,
203    view: &crate::PhysicalRowView<'_>,
204) -> ExecResult<Value> {
205    match source {
206        JoinOutputSource::Input(position) => {
207            Ok(view.value_at(*position).unwrap_or(&Value::Null).clone())
208        }
209        JoinOutputSource::Cast { input, ty } => cast_value(
210            view.value_at(*input).unwrap_or(&Value::Null),
211            &ty.sql_name(),
212        )
213        .map_err(ExecError::from),
214        JoinOutputSource::Coalesce { left, right, ty } => {
215            let target = ty.sql_name();
216            let left = cast_value(view.value_at(*left).unwrap_or(&Value::Null), &target)?;
217            if !matches!(left, Value::Null) {
218                return Ok(left);
219            }
220            cast_value(view.value_at(*right).unwrap_or(&Value::Null), &target)
221                .map_err(ExecError::from)
222        }
223    }
224}
225
226#[cfg(test)]
227mod tests {
228    use std::collections::BTreeMap;
229
230    use super::*;
231    use crate::physical::run_to_rows;
232    use crate::TableScan;
233    use uqa_sql::expr::RowLookup;
234
235    #[test]
236    fn schema_only_merge_reuses_input_slots_and_keeps_qualified_aliases() {
237        let child = TableScan::from_rows(
238            vec![
239                "l.id".into(),
240                "l.name".into(),
241                "r.id".into(),
242                "r.note".into(),
243            ],
244            vec![BTreeMap::from([
245                ("l.id".into(), Value::Int(1)),
246                ("l.name".into(), Value::Str("left".into())),
247                ("r.id".into(), Value::Int(1)),
248                ("r.note".into(), Value::Str("right".into())),
249            ])],
250        );
251        let mut output = JoinOutput::try_new(
252            Box::new(child),
253            vec![
254                (
255                    "id".into(),
256                    ColumnIdentity::unqualified("id"),
257                    JoinOutputSource::Input(0),
258                ),
259                (
260                    "name".into(),
261                    ColumnIdentity::qualified("l", "name"),
262                    JoinOutputSource::Input(1),
263                ),
264                (
265                    "note".into(),
266                    ColumnIdentity::qualified("r", "note"),
267                    JoinOutputSource::Input(3),
268                ),
269            ],
270            vec![
271                (
272                    ColumnIdentity::qualified("l", "id"),
273                    JoinOutputSource::Input(0),
274                ),
275                (
276                    ColumnIdentity::qualified("r", "id"),
277                    JoinOutputSource::Input(2),
278                ),
279            ],
280        )
281        .unwrap();
282        output.open().unwrap();
283        let batch = output.next().unwrap().unwrap();
284        assert_eq!(batch.schema.columns(), ["id", "name", "note"]);
285        let view = batch.schema.view(&batch.rows[0]);
286        assert_eq!(view.qualified_column("l", "id"), Some(&Value::Int(1)));
287        assert_eq!(view.qualified_column("r", "id"), Some(&Value::Int(1)));
288        output.close().unwrap();
289    }
290
291    #[test]
292    fn full_merge_coalesces_only_the_requested_column() {
293        let child = TableScan::from_rows(
294            vec!["l.id".into(), "r.id".into()],
295            vec![BTreeMap::from([
296                ("l.id".into(), Value::Null),
297                ("r.id".into(), Value::Int(3)),
298            ])],
299        );
300        let mut output = JoinOutput::try_new(
301            Box::new(child),
302            vec![(
303                "id".into(),
304                ColumnIdentity::unqualified("id"),
305                JoinOutputSource::Coalesce {
306                    left: 0,
307                    right: 1,
308                    ty: ColumnType::Integer,
309                },
310            )],
311            Vec::new(),
312        )
313        .unwrap();
314        let (_, rows) = run_to_rows(&mut output).unwrap();
315        assert_eq!(rows[0]["id"], Value::Int(3));
316    }
317}
uqa_execution/join_output.rs

uqa_execution/
join_output.rs
