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    /// Physical ordering after the selected positions have been remapped.
47    /// Executor-only projection attributes are addressed structurally and
48    /// never enter the public SQL column namespace.
49    ordering: Vec<PhysicalOrder>,
50    rebind_lock_qualifier: Option<Arc<str>>,
51    discard_lock_origins: bool,
52    compact_slots: Option<Vec<Option<usize>>>,
53}
54
55impl<'a> ColumnSelection<'a> {
56    /// Keep the visible layout unchanged while retiring consumed executor-only
57    /// attributes from the schema namespace.
58    pub fn dropping_internal_attributes(
59        child: Box<dyn PhysicalOperator + 'a>,
60        columns: &[uqa_sql::ast::InternalColumnRef],
61    ) -> Self {
62        let schema = RowSchema::without_internal_attributes(child.row_schema(), columns);
63        let ordering = child.output_ordering().to_vec();
64        Self {
65            child,
66            schema,
67            ordering,
68            rebind_lock_qualifier: None,
69            discard_lock_origins: false,
70            compact_slots: None,
71        }
72    }
73
74    pub fn new(child: Box<dyn PhysicalOperator + 'a>, columns: Vec<String>) -> Self {
75        let columns = columns
76            .into_iter()
77            .map(|column| (column.clone(), column))
78            .collect();
79        Self::with_mapping(child, columns)
80    }
81
82    pub fn with_mapping(
83        child: Box<dyn PhysicalOperator + 'a>,
84        columns: Vec<(String, String)>,
85    ) -> Self {
86        let input_positions = columns
87            .iter()
88            .map(|(_, input)| child.row_schema().position(input))
89            .collect::<Vec<_>>();
90        let ordering = remap_ordering(child.output_ordering(), &input_positions);
91        let schema = RowSchema::select(child.row_schema(), &columns);
92        Self {
93            child,
94            schema,
95            ordering,
96            rebind_lock_qualifier: None,
97            discard_lock_origins: false,
98            compact_slots: None,
99        }
100    }
101
102    /// Select and rename logical input positions without resolving them by
103    /// name. SQL result shaping uses this when repeated public labels must
104    /// remain distinct even though their qualified input names differ.
105    pub fn with_positions(
106        child: Box<dyn PhysicalOperator + 'a>,
107        columns: Vec<(String, usize)>,
108    ) -> Self {
109        let input_positions = columns
110            .iter()
111            .map(|(_, position)| Some(*position))
112            .collect::<Vec<_>>();
113        let ordering = remap_ordering(child.output_ordering(), &input_positions);
114        let schema = RowSchema::remap_positions(child.row_schema(), &columns, &[]);
115        Self {
116            child,
117            schema,
118            ordering,
119            rebind_lock_qualifier: None,
120            discard_lock_origins: false,
121            compact_slots: None,
122        }
123    }
124
125    /// Select flattened physical slots, including executor-only internal
126    /// attributes that deliberately have no SQL-visible logical position.
127    pub fn with_physical_positions(
128        child: Box<dyn PhysicalOperator + 'a>,
129        columns: Vec<(String, usize)>,
130    ) -> Self {
131        let input_positions = columns
132            .iter()
133            .map(|(_, physical)| {
134                (0..child.row_schema().len())
135                    .find(|logical| child.row_schema().physical_slot(*logical) == Some(*physical))
136            })
137            .collect::<Vec<_>>();
138        let ordering = remap_ordering(child.output_ordering(), &input_positions);
139        let columns = columns
140            .into_iter()
141            .map(|(label, physical)| {
142                let ty = child.row_schema().physical_type(physical).cloned();
143                (
144                    label.clone(),
145                    ColumnIdentity::unqualified(label),
146                    physical,
147                    ty,
148                )
149            })
150            .collect::<Vec<_>>();
151        let schema = RowSchema::remap_typed_physical_identities(child.row_schema(), &columns, &[]);
152        Self {
153            child,
154            schema,
155            ordering,
156            rebind_lock_qualifier: None,
157            discard_lock_origins: false,
158            compact_slots: None,
159        }
160    }
161
162    /// Select logical input positions and assign explicit structured SQL identities without encoding qualifiers into output labels.
163    pub fn with_identities(
164        child: Box<dyn PhysicalOperator + 'a>,
165        columns: Vec<(String, ColumnIdentity, usize)>,
166    ) -> Self {
167        let input_positions = columns
168            .iter()
169            .map(|(_, _, position)| Some(*position))
170            .collect::<Vec<_>>();
171        let ordering = remap_ordering(child.output_ordering(), &input_positions);
172        let columns = columns
173            .into_iter()
174            .map(|(label, identity, position)| {
175                let ty = child.row_schema().column_type(position).cloned();
176                (label, identity, position, ty)
177            })
178            .collect::<Vec<_>>();
179        let schema = RowSchema::remap_typed_identities(child.row_schema(), &columns, &[]);
180        Self {
181            child,
182            schema,
183            ordering,
184            rebind_lock_qualifier: None,
185            discard_lock_origins: false,
186            compact_slots: None,
187        }
188    }
189
190    /// Hide recursive-CTE generated columns from `*` while retaining them as real qualified and unqualified lookup aliases. `PostgreSQL` adds SEARCH/CYCLE state after expanding the recursive term's wildcard, but still permits the recursive term to name the generated cycle columns explicitly.
191    pub fn hiding_trailing_columns(
192        child: Box<dyn PhysicalOperator + 'a>,
193        visible: usize,
194        qualifier: &str,
195    ) -> Self {
196        let visible = visible.min(child.row_schema().len());
197        let input_positions = (0..visible).map(Some).collect::<Vec<_>>();
198        let ordering = remap_ordering(child.output_ordering(), &input_positions);
199        let columns = (0..visible)
200            .map(|position| {
201                let label = child.row_schema().columns()[position].clone();
202                let identity = if qualifier.is_empty() {
203                    ColumnIdentity::unqualified(&label)
204                } else {
205                    ColumnIdentity::qualified(qualifier, &label)
206                };
207                let ty = child.row_schema().column_type(position).cloned();
208                (label, identity, position, ty)
209            })
210            .collect::<Vec<_>>();
211        let mut aliases = Vec::new();
212        for position in visible..child.row_schema().len() {
213            let name = child.row_schema().columns()[position].clone();
214            aliases.push((ColumnIdentity::unqualified(&name), position));
215            if !qualifier.is_empty() {
216                aliases.push((ColumnIdentity::qualified(qualifier, name), position));
217            }
218        }
219        let schema = RowSchema::remap_typed_identities(child.row_schema(), &columns, &aliases);
220        Self {
221            child,
222            schema,
223            ordering,
224            rebind_lock_qualifier: None,
225            discard_lock_origins: true,
226            compact_slots: None,
227        }
228    }
229
230    /// Select positions under identities that replace, rather than extend, the child's SQL namespace.
231    pub fn with_fresh_identities(
232        child: Box<dyn PhysicalOperator + 'a>,
233        columns: Vec<(String, ColumnIdentity, usize)>,
234    ) -> Self {
235        let input_positions = columns
236            .iter()
237            .map(|(_, _, position)| Some(*position))
238            .collect::<Vec<_>>();
239        let ordering = remap_ordering(child.output_ordering(), &input_positions);
240        let columns = columns
241            .into_iter()
242            .map(|(label, identity, position)| {
243                let ty = child.row_schema().column_type(position).cloned();
244                (label, identity, position, ty)
245            })
246            .collect::<Vec<_>>();
247        let schema =
248            RowSchema::remap_typed_identities_without_input_aliases(child.row_schema(), &columns);
249        Self {
250            child,
251            schema,
252            ordering,
253            rebind_lock_qualifier: None,
254            discard_lock_origins: false,
255            compact_slots: None,
256        }
257    }
258
259    /// 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.
260    pub fn compacting_with_positions(
261        child: Box<dyn PhysicalOperator + 'a>,
262        columns: Vec<(String, usize)>,
263    ) -> Self {
264        let input_positions = columns
265            .iter()
266            .map(|(_, position)| Some(*position))
267            .collect::<Vec<_>>();
268        let ordering = remap_ordering(child.output_ordering(), &input_positions);
269        let selected = RowSchema::remap_positions(child.row_schema(), &columns, &[]);
270        let (schema, source_slots) = selected.canonical_projection();
271        let identity_layout = child.row_schema().physical_width() == source_slots.len()
272            && source_slots
273                .iter()
274                .enumerate()
275                .all(|(position, slot)| *slot == position);
276        Self {
277            child,
278            schema,
279            ordering,
280            rebind_lock_qualifier: None,
281            discard_lock_origins: true,
282            compact_slots: (!identity_layout).then(|| source_slots.into_iter().map(Some).collect()),
283        }
284    }
285
286    /// Attribute inner lock origins to this source qualifier so `FOR UPDATE OF` a view, CTE, or subquery does not lock sibling join inputs.
287    #[must_use]
288    pub fn rebinding_lock_origins(mut self, qualifier: impl Into<String>) -> Self {
289        let qualifier = qualifier.into();
290        if !qualifier.is_empty() {
291            self.rebind_lock_qualifier = Some(Arc::from(qualifier));
292        }
293        self
294    }
295
296    /// Attribute every structurally carried retrieval score to this relation alias without introducing a hidden SQL column name.
297    #[must_use]
298    pub fn rebinding_score_sources(mut self, qualifier: impl Into<String>) -> Self {
299        let qualifier = qualifier.into();
300        self.schema = RowSchema::with_rebound_score_sources(
301            &self.schema,
302            (!qualifier.is_empty()).then_some(qualifier.as_str()),
303        );
304        self
305    }
306
307    /// Remove row-lock identities at a relational row-identity barrier.
308    #[must_use]
309    pub fn discarding_lock_origins(mut self) -> Self {
310        self.rebind_lock_qualifier = None;
311        self.discard_lock_origins = true;
312        self
313    }
314
315    fn select_batch(&self, batch: Batch) -> Batch {
316        let mut rows = match self.compact_slots.as_ref() {
317            Some(slots) => batch
318                .rows
319                .into_iter()
320                .map(|row| {
321                    row.project_with_values(slots.iter().map(|slot| {
322                        slot.map_or(
323                            RowProjectionValue::Owned(Value::Null),
324                            RowProjectionValue::InputSlot,
325                        )
326                    }))
327                })
328                .collect(),
329            None => batch.rows,
330        };
331        if self.discard_lock_origins {
332            for row in &mut rows {
333                row.discard_lock_origins_mut();
334            }
335        } else if let Some(qualifier) = self.rebind_lock_qualifier.as_ref() {
336            for row in &mut rows {
337                row.rebind_lock_origin_qualifiers_mut(Arc::clone(qualifier));
338            }
339        }
340        Batch::from_physical_rows(self.schema.clone(), rows)
341    }
342}
343
344impl PhysicalOperator for ColumnSelection<'_> {
345    fn row_schema(&self) -> &RowSchema {
346        &self.schema
347    }
348
349    fn estimated_cardinality(&self) -> Option<u64> {
350        self.child.estimated_cardinality()
351    }
352
353    fn output_ordering(&self) -> &[PhysicalOrder] {
354        &self.ordering
355    }
356
357    fn backward_scan_support(&self) -> crate::BackwardScanSupport {
358        self.child.backward_scan_support()
359    }
360
361    fn open(&mut self) -> ExecResult<()> {
362        self.child.open()
363    }
364
365    fn next(&mut self) -> ExecResult<Option<Batch>> {
366        let Some(batch) = self.child.next()? else {
367            return Ok(None);
368        };
369        Ok(Some(self.select_batch(batch)))
370    }
371
372    fn next_direction(
373        &mut self,
374        direction: crate::PhysicalScanDirection,
375    ) -> ExecResult<Option<Batch>> {
376        let Some(batch) = self.child.next_direction(direction)? else {
377            return Ok(None);
378        };
379        Ok(Some(self.select_batch(batch)))
380    }
381
382    fn rewind(&mut self) -> ExecResult<()> {
383        self.child.rewind()
384    }
385
386    fn close(&mut self) -> ExecResult<()> {
387        self.child.close()
388    }
389}
390
391#[cfg(test)]
392mod tests {
393    use std::collections::BTreeMap;
394
395    use uqa_core::Value;
396
397    use super::*;
398    use crate::physical::run_to_rows;
399    use crate::scan::TableScan;
400
401    #[test]
402    fn selects_computed_columns_without_leaking_sort_inputs() {
403        let row = BTreeMap::from([
404            ("source".to_string(), Value::Int(1)),
405            ("alias".to_string(), Value::Int(2)),
406        ]);
407        let scan = TableScan::from_rows(vec!["source".into(), "alias".into()], vec![row]);
408        let mut selection = ColumnSelection::new(Box::new(scan), vec!["alias".into()]);
409        let (schema, rows) = run_to_rows(&mut selection).unwrap();
410        assert_eq!(schema, vec!["alias"]);
411        assert_eq!(rows[0], BTreeMap::from([("alias".into(), Value::Int(2))]));
412    }
413
414    #[test]
415    fn renames_collision_free_physical_columns() {
416        let row = BTreeMap::from([
417            ("source".to_string(), Value::Int(1)),
418            ("__projection_0".to_string(), Value::Int(2)),
419        ]);
420        let scan = TableScan::from_rows(vec!["source".into(), "__projection_0".into()], vec![row]);
421        let mut selection = ColumnSelection::with_mapping(
422            Box::new(scan),
423            vec![("source".into(), "__projection_0".into())],
424        );
425        let (schema, rows) = run_to_rows(&mut selection).unwrap();
426        assert_eq!(schema, vec!["source"]);
427        assert_eq!(rows[0], BTreeMap::from([("source".into(), Value::Int(2))]));
428    }
429
430    #[test]
431    fn renames_repeated_columns_by_position() {
432        let scan = TableScan::from_rows(
433            vec!["left.value".into(), "right.value".into()],
434            vec![BTreeMap::from([
435                ("left.value".into(), Value::Int(1)),
436                ("right.value".into(), Value::Int(2)),
437            ])],
438        );
439        let mut selection = ColumnSelection::with_positions(
440            Box::new(scan),
441            vec![("value".into(), 0), ("value".into(), 1)],
442        );
443        let batches = crate::physical::run_to_batches(&mut selection).unwrap();
444        assert_eq!(batches[0].schema.columns(), ["value", "value"]);
445        let row = batches[0].schema.view(&batches[0].rows[0]);
446        assert_eq!(row.value_at(0), Some(&Value::Int(1)));
447        assert_eq!(row.value_at(1), Some(&Value::Int(2)));
448    }
449
450    #[test]
451    fn ordering_is_remapped_by_selected_position() {
452        let ordering = vec![PhysicalOrder {
453            position: 2,
454            descending: false,
455            nulls_first: None,
456            nullable: false,
457        }];
458        let remapped = remap_ordering(&ordering, &[Some(2), Some(0)]);
459        assert_eq!(remapped[0].position, 0);
460        assert!(remap_ordering(&ordering, &[Some(0), Some(1)]).is_empty());
461    }
462
463    #[test]
464    fn row_identity_barrier_discards_lock_origins_in_place() {
465        let schema = RowSchema::new(vec!["id".into()]);
466        let row = crate::PhysicalRow::from_values(vec![Value::Int(1)])
467            .with_lock_origin(crate::RowLockOrigin::new("accounts", "public.accounts", 1));
468        let scan = TableScan::from_physical_rows(schema, vec![row]);
469        let mut barrier = ColumnSelection::with_positions(Box::new(scan), vec![("id".into(), 0)])
470            .discarding_lock_origins();
471
472        let batches = crate::physical::run_to_batches(&mut barrier).unwrap();
473        assert!(batches[0].rows[0].lock_origins().is_empty());
474    }
475
476    #[test]
477    fn explicit_compaction_canonicalizes_a_wider_physical_layout() {
478        let source = RowSchema::new(vec!["unused".into(), "value".into()]);
479        let selected = RowSchema::select(&source, &[("value".into(), "value".into())]);
480        let row = crate::PhysicalRow::from_values(vec![Value::Int(1), Value::Int(7)]);
481        let scan = TableScan::from_physical_rows(selected, vec![row]);
482        let mut compact =
483            ColumnSelection::compacting_with_positions(Box::new(scan), vec![("value".into(), 0)]);
484
485        compact.open().unwrap();
486        let batch = compact.next().unwrap().unwrap();
487        compact.close().unwrap();
488
489        assert_eq!(batch.schema.physical_width(), 1);
490        assert_eq!(
491            batch.schema.view(&batch.rows[0]).get("value"),
492            Some(&Value::Int(7))
493        );
494    }
495}
uqa_execution/column_selection.rs

uqa_execution/
column_selection.rs
