uqa_execution/
column_selection.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Schema-only projection used after an expression-producing stage.
8
9use std::sync::Arc;
10
11use uqa_core::Value;
12
13use crate::{
14    Batch, ColumnIdentity, ExecResult, PhysicalOperator, PhysicalOrder, RowProjectionValue,
15    RowSchema,
16};
17
18fn remap_ordering(
19    ordering: &[PhysicalOrder],
20    input_positions: &[Option<usize>],
21) -> Vec<PhysicalOrder> {
22    ordering
23        .iter()
24        .map_while(|order| {
25            let position = input_positions
26                .iter()
27                .position(|input| *input == Some(order.position))?;
28            Some(PhysicalOrder {
29                position,
30                ..order.clone()
31            })
32        })
33        .collect()
34}
35
36/// Select already-computed columns without evaluating expressions again.
37///
38/// This is intentionally distinct from [`crate::Project`]: SQL `ORDER BY`
39/// may reference both source columns and SELECT aliases, so the expression
40/// projection first appends aliases, Sort consumes that augmented row, and
41/// `ColumnSelection` removes the non-output source columns afterwards. Volatile
42/// projection expressions are therefore evaluated exactly once.
43pub struct ColumnSelection<'a> {
44    child: Box<dyn PhysicalOperator + 'a>,
45    schema: RowSchema,
46    /// `(output_name, input_name)` pairs. Keeping the physical input name
47    /// separate lets a preceding projection expose SELECT-list expressions
48    /// under collision-free internal names while this final, non-evaluating
49    /// operator restores the public SQL column names.
50    ordering: Vec<PhysicalOrder>,
51    rebind_lock_qualifier: Option<Arc<str>>,
52    discard_lock_origins: bool,
53    compact_slots: Option<Vec<Option<usize>>>,
54}
55
56impl<'a> ColumnSelection<'a> {
57    pub fn new(child: Box<dyn PhysicalOperator + 'a>, columns: Vec<String>) -> Self {
58        let columns = columns
59            .into_iter()
60            .map(|column| (column.clone(), column))
61            .collect();
62        Self::with_mapping(child, columns)
63    }
64
65    pub fn with_mapping(
66        child: Box<dyn PhysicalOperator + 'a>,
67        columns: Vec<(String, String)>,
68    ) -> Self {
69        let input_positions = columns
70            .iter()
71            .map(|(_, input)| child.row_schema().position(input))
72            .collect::<Vec<_>>();
73        let ordering = remap_ordering(child.output_ordering(), &input_positions);
74        let schema = RowSchema::select(child.row_schema(), &columns);
75        Self {
76            child,
77            schema,
78            ordering,
79            rebind_lock_qualifier: None,
80            discard_lock_origins: false,
81            compact_slots: None,
82        }
83    }
84
85    /// Select and rename logical input positions without resolving them by
86    /// name. SQL result shaping uses this when repeated public labels must
87    /// remain distinct even though their qualified input names differ.
88    pub fn with_positions(
89        child: Box<dyn PhysicalOperator + 'a>,
90        columns: Vec<(String, usize)>,
91    ) -> Self {
92        let input_positions = columns
93            .iter()
94            .map(|(_, position)| Some(*position))
95            .collect::<Vec<_>>();
96        let ordering = remap_ordering(child.output_ordering(), &input_positions);
97        let schema = RowSchema::remap_positions(child.row_schema(), &columns, &[]);
98        Self {
99            child,
100            schema,
101            ordering,
102            rebind_lock_qualifier: None,
103            discard_lock_origins: false,
104            compact_slots: None,
105        }
106    }
107
108    /// Select logical input positions and assign explicit structured SQL identities without encoding qualifiers into output labels.
109    pub fn with_identities(
110        child: Box<dyn PhysicalOperator + 'a>,
111        columns: Vec<(String, ColumnIdentity, usize)>,
112    ) -> Self {
113        let input_positions = columns
114            .iter()
115            .map(|(_, _, position)| Some(*position))
116            .collect::<Vec<_>>();
117        let ordering = remap_ordering(child.output_ordering(), &input_positions);
118        let columns = columns
119            .into_iter()
120            .map(|(label, identity, position)| {
121                let ty = child.row_schema().column_type(position).cloned();
122                (label, identity, position, ty)
123            })
124            .collect::<Vec<_>>();
125        let schema = RowSchema::remap_typed_identities(child.row_schema(), &columns, &[]);
126        Self {
127            child,
128            schema,
129            ordering,
130            rebind_lock_qualifier: None,
131            discard_lock_origins: false,
132            compact_slots: None,
133        }
134    }
135
136    /// Select logical input positions and compact them to one canonical positional layout. Use this only at an explicit state boundary, such as recursive working-table materialization, where independently planned inputs must share an identical physical schema.
137    pub fn compacting_with_positions(
138        child: Box<dyn PhysicalOperator + 'a>,
139        columns: Vec<(String, usize)>,
140    ) -> Self {
141        let input_positions = columns
142            .iter()
143            .map(|(_, position)| Some(*position))
144            .collect::<Vec<_>>();
145        let ordering = remap_ordering(child.output_ordering(), &input_positions);
146        let names = columns
147            .iter()
148            .map(|(name, _)| name.clone())
149            .collect::<Vec<_>>();
150        let types = columns
151            .iter()
152            .map(|(_, position)| child.row_schema().column_type(*position).cloned())
153            .collect::<Vec<_>>();
154        let slots = columns
155            .iter()
156            .map(|(_, position)| child.row_schema().physical_slot(*position))
157            .collect::<Vec<_>>();
158        let identity_layout = child.row_schema().physical_width() == slots.len()
159            && slots
160                .iter()
161                .enumerate()
162                .all(|(position, slot)| *slot == Some(position));
163        Self {
164            child,
165            schema: RowSchema::with_types(names, types),
166            ordering,
167            rebind_lock_qualifier: None,
168            discard_lock_origins: true,
169            compact_slots: (!identity_layout).then_some(slots),
170        }
171    }
172
173    /// Attribute inner lock origins to this source qualifier so `FOR UPDATE OF` a view, CTE, or subquery does not lock sibling join inputs.
174    #[must_use]
175    pub fn rebinding_lock_origins(mut self, qualifier: impl Into<String>) -> Self {
176        let qualifier = qualifier.into();
177        if !qualifier.is_empty() {
178            self.rebind_lock_qualifier = Some(Arc::from(qualifier));
179        }
180        self
181    }
182
183    /// Remove row-lock identities at a relational row-identity barrier.
184    #[must_use]
185    pub fn discarding_lock_origins(mut self) -> Self {
186        self.rebind_lock_qualifier = None;
187        self.discard_lock_origins = true;
188        self
189    }
190}
191
192impl PhysicalOperator for ColumnSelection<'_> {
193    fn row_schema(&self) -> &RowSchema {
194        &self.schema
195    }
196
197    fn estimated_cardinality(&self) -> Option<u64> {
198        self.child.estimated_cardinality()
199    }
200
201    fn output_ordering(&self) -> &[PhysicalOrder] {
202        &self.ordering
203    }
204
205    fn open(&mut self) -> ExecResult<()> {
206        self.child.open()
207    }
208
209    fn next(&mut self) -> ExecResult<Option<Batch>> {
210        let Some(batch) = self.child.next()? else {
211            return Ok(None);
212        };
213        let mut rows = match self.compact_slots.as_ref() {
214            Some(slots) => batch
215                .rows
216                .into_iter()
217                .map(|row| {
218                    row.project_with_values(slots.iter().map(|slot| {
219                        slot.map_or(
220                            RowProjectionValue::Owned(Value::Null),
221                            RowProjectionValue::InputSlot,
222                        )
223                    }))
224                })
225                .collect(),
226            None => batch.rows,
227        };
228        if self.discard_lock_origins {
229            for row in &mut rows {
230                row.discard_lock_origins_mut();
231            }
232        } else if let Some(qualifier) = self.rebind_lock_qualifier.as_ref() {
233            for row in &mut rows {
234                row.rebind_lock_origin_qualifiers_mut(Arc::clone(qualifier));
235            }
236        }
237        Ok(Some(Batch::from_physical_rows(self.schema.clone(), rows)))
238    }
239
240    fn close(&mut self) -> ExecResult<()> {
241        self.child.close()
242    }
243}
244
245#[cfg(test)]
246mod tests {
247    use std::collections::BTreeMap;
248
249    use uqa_core::Value;
250
251    use super::*;
252    use crate::physical::run_to_rows;
253    use crate::scan::TableScan;
254
255    #[test]
256    fn selects_computed_columns_without_leaking_sort_inputs() {
257        let row = BTreeMap::from([
258            ("source".to_string(), Value::Int(1)),
259            ("alias".to_string(), Value::Int(2)),
260        ]);
261        let scan = TableScan::from_rows(vec!["source".into(), "alias".into()], vec![row]);
262        let mut selection = ColumnSelection::new(Box::new(scan), vec!["alias".into()]);
263        let (schema, rows) = run_to_rows(&mut selection).unwrap();
264        assert_eq!(schema, vec!["alias"]);
265        assert_eq!(rows[0], BTreeMap::from([("alias".into(), Value::Int(2))]));
266    }
267
268    #[test]
269    fn renames_collision_free_physical_columns() {
270        let row = BTreeMap::from([
271            ("source".to_string(), Value::Int(1)),
272            ("__projection_0".to_string(), Value::Int(2)),
273        ]);
274        let scan = TableScan::from_rows(vec!["source".into(), "__projection_0".into()], vec![row]);
275        let mut selection = ColumnSelection::with_mapping(
276            Box::new(scan),
277            vec![("source".into(), "__projection_0".into())],
278        );
279        let (schema, rows) = run_to_rows(&mut selection).unwrap();
280        assert_eq!(schema, vec!["source"]);
281        assert_eq!(rows[0], BTreeMap::from([("source".into(), Value::Int(2))]));
282    }
283
284    #[test]
285    fn renames_repeated_columns_by_position() {
286        let scan = TableScan::from_rows(
287            vec!["left.value".into(), "right.value".into()],
288            vec![BTreeMap::from([
289                ("left.value".into(), Value::Int(1)),
290                ("right.value".into(), Value::Int(2)),
291            ])],
292        );
293        let mut selection = ColumnSelection::with_positions(
294            Box::new(scan),
295            vec![("value".into(), 0), ("value".into(), 1)],
296        );
297        let batches = crate::physical::run_to_batches(&mut selection).unwrap();
298        assert_eq!(batches[0].schema.columns(), ["value", "value"]);
299        let row = batches[0].schema.view(&batches[0].rows[0]);
300        assert_eq!(row.value_at(0), Some(&Value::Int(1)));
301        assert_eq!(row.value_at(1), Some(&Value::Int(2)));
302    }
303
304    #[test]
305    fn ordering_is_remapped_by_selected_position() {
306        let ordering = vec![PhysicalOrder {
307            position: 2,
308            descending: false,
309            nulls_first: None,
310            nullable: false,
311        }];
312        let remapped = remap_ordering(&ordering, &[Some(2), Some(0)]);
313        assert_eq!(remapped[0].position, 0);
314        assert!(remap_ordering(&ordering, &[Some(0), Some(1)]).is_empty());
315    }
316
317    #[test]
318    fn row_identity_barrier_discards_lock_origins_in_place() {
319        let schema = RowSchema::new(vec!["id".into()]);
320        let row = crate::PhysicalRow::from_values(vec![Value::Int(1)])
321            .with_lock_origin(crate::RowLockOrigin::new("accounts", "public.accounts", 1));
322        let scan = TableScan::from_physical_rows(schema, vec![row]);
323        let mut barrier = ColumnSelection::with_positions(Box::new(scan), vec![("id".into(), 0)])
324            .discarding_lock_origins();
325
326        let batches = crate::physical::run_to_batches(&mut barrier).unwrap();
327        assert!(batches[0].rows[0].lock_origins().is_empty());
328    }
329
330    #[test]
331    fn explicit_compaction_canonicalizes_a_wider_physical_layout() {
332        let source = RowSchema::new(vec!["unused".into(), "value".into()]);
333        let selected = RowSchema::select(&source, &[("value".into(), "value".into())]);
334        let row = crate::PhysicalRow::from_values(vec![Value::Int(1), Value::Int(7)]);
335        let scan = TableScan::from_physical_rows(selected, vec![row]);
336        let mut compact =
337            ColumnSelection::compacting_with_positions(Box::new(scan), vec![("value".into(), 0)]);
338
339        compact.open().unwrap();
340        let batch = compact.next().unwrap().unwrap();
341        compact.close().unwrap();
342
343        assert_eq!(batch.schema.physical_width(), 1);
344        assert_eq!(
345            batch.schema.view(&batch.rows[0]).get("value"),
346            Some(&Value::Int(7))
347        );
348    }
349}
uqa_execution/column_selection.rs

uqa_execution/
column_selection.rs
