Skip to main content

uqa_execution/batch/
schema_projection.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Projection, canonicalization, and position remapping for row schemas.
8
9use super::{
10    ColumnIdentity, ColumnType, ExecError, ExecResult, HashMap, PhysicalLayout, ProjectedSlot,
11    RowSchema, SchemaBuildMetadata, ScoreSource, NULL_SLOT,
12};
13
14impl RowSchema {
15    /// Remove selected executor-only identities after their consumer has run.
16    /// The physical fragments remain shareable until the next canonical
17    /// boundary, where now-unreferenced slots are naturally discarded.
18    pub fn without_internal_attributes(
19        input: &Self,
20        columns: &[uqa_sql::ast::InternalColumnRef],
21    ) -> Self {
22        let mut internal = input.index.executor_attributes.clone();
23        let mut internal_types = input.index.cold.executor_attribute_types.clone();
24        for column in columns {
25            internal.remove(column);
26            internal_types.remove(column);
27        }
28        let score_sources = input
29            .index
30            .cold
31            .score_sources
32            .iter()
33            .filter(|source| internal.contains_key(&source.column))
34            .cloned()
35            .collect();
36        Self::from_typed_parts_with_aliases_and_exact_precedence(
37            input.columns().to_vec(),
38            input.identities().to_vec(),
39            input.column_types().to_vec(),
40            input.index.slots.to_vec(),
41            input.physical_width(),
42            SchemaBuildMetadata {
43                aliases: input.index.aliases.clone(),
44                alias_types: input.index.cold.aliases.clone(),
45                internal,
46                internal_types,
47                score_sources,
48                wildcard_hidden: input.index.cold.wildcard_hidden.clone(),
49                binding_only: input.index.cold.binding_only.clone(),
50                ..SchemaBuildMetadata::default()
51            },
52        )
53    }
54
55    /// Select and optionally rename logical columns while retaining the
56    /// child's physical fragments.
57    pub fn select(input: &Self, columns: &[(String, String)]) -> Self {
58        let output_names = columns
59            .iter()
60            .map(|(output, _)| output.clone())
61            .collect::<Vec<_>>();
62        let slots = columns
63            .iter()
64            .map(|(_, source)| input.exact_slot(source).unwrap_or(NULL_SLOT))
65            .collect();
66        let types = columns
67            .iter()
68            .map(|(_, source)| input.exact_type(source).cloned())
69            .collect();
70        let identities = output_names
71            .iter()
72            .cloned()
73            .map(ColumnIdentity::unqualified)
74            .collect();
75        Self::from_typed_parts_with_aliases_and_exact_precedence(
76            output_names,
77            identities,
78            types,
79            slots,
80            input.physical_width(),
81            SchemaBuildMetadata {
82                aliases: HashMap::new(),
83                alias_types: HashMap::new(),
84                internal: input.index.executor_attributes.clone(),
85                internal_types: input.index.cold.executor_attribute_types.clone(),
86                score_sources: input.index.cold.score_sources.clone(),
87                binding_only: HashMap::new(),
88                ..SchemaBuildMetadata::default()
89            },
90        )
91    }
92
93    /// Build a scalar-projection schema without rebuilding direct input values. A non-pass-through projection hides child identities logically while retaining their physical fragments; an appending projection preserves the child schema and replaces duplicate labels with `PostgreSQL` map-insertion semantics.
94    pub(crate) fn project_with_sources(
95        input: &Self,
96        projected: Vec<(String, Option<ColumnType>, ProjectedSlot)>,
97        projected_internal: Vec<(
98            uqa_sql::ast::InternalColumnRef,
99            Option<ColumnType>,
100            ProjectedSlot,
101        )>,
102        computed_count: usize,
103        pass_through: bool,
104    ) -> Self {
105        let resolve_slot = |source: ProjectedSlot| match source {
106            ProjectedSlot::Input(slot) => slot.unwrap_or(NULL_SLOT),
107            ProjectedSlot::Computed(position) => input.physical_width() + position,
108        };
109        let physical_width = input.physical_width() + computed_count;
110        let mut internal = input.index.executor_attributes.clone();
111        let mut internal_types = input.index.cold.executor_attribute_types.clone();
112        for (column, ty, source) in projected_internal {
113            internal.insert(column, resolve_slot(source));
114            internal_types.insert(column, ty);
115        }
116
117        if pass_through {
118            let mut columns = input.columns().to_vec();
119            let mut identities = input.identities().to_vec();
120            let mut types = input.column_types().to_vec();
121            let mut slots = input.index.slots.to_vec();
122            let mut wildcard_hidden = input.index.cold.wildcard_hidden.clone();
123            for (name, ty, source) in projected {
124                let slot = resolve_slot(source);
125                if let Some(position) = columns.iter().position(|column| column == &name) {
126                    slots[position] = slot;
127                    identities[position] = ColumnIdentity::unqualified(name);
128                    types[position] = ty;
129                    wildcard_hidden.remove(&position);
130                } else {
131                    identities.push(ColumnIdentity::unqualified(name.clone()));
132                    columns.push(name);
133                    types.push(ty);
134                    slots.push(slot);
135                }
136            }
137            return Self::from_typed_parts_with_aliases_and_exact_precedence(
138                columns,
139                identities,
140                types,
141                slots,
142                physical_width,
143                SchemaBuildMetadata {
144                    aliases: input.index.aliases.clone(),
145                    alias_types: input.index.cold.aliases.clone(),
146                    internal,
147                    internal_types,
148                    score_sources: input.index.cold.score_sources.clone(),
149                    wildcard_hidden,
150                    binding_only: input.index.cold.binding_only.clone(),
151                    ..SchemaBuildMetadata::default()
152                },
153            );
154        }
155
156        let mut columns = Vec::with_capacity(projected.len());
157        let mut identities = Vec::with_capacity(projected.len());
158        let mut types = Vec::with_capacity(projected.len());
159        let mut slots = Vec::with_capacity(projected.len());
160        for (name, ty, source) in projected {
161            slots.push(resolve_slot(source));
162            identities.push(ColumnIdentity::unqualified(name.clone()));
163            columns.push(name);
164            types.push(ty);
165        }
166        Self::from_typed_parts_with_aliases_and_exact_precedence(
167            columns,
168            identities,
169            types,
170            slots,
171            physical_width,
172            SchemaBuildMetadata {
173                aliases: HashMap::new(),
174                alias_types: HashMap::new(),
175                internal,
176                internal_types,
177                score_sources: input.index.cold.score_sources.clone(),
178                binding_only: HashMap::new(),
179                ..SchemaBuildMetadata::default()
180            },
181        )
182    }
183
184    /// Build a compact positional layout for a blocking or spill boundary.
185    /// Logical columns and hidden lookup aliases are remapped to a deduplicated
186    /// list of referenced physical slots; projecting a row through the returned
187    /// slot list shares its existing value fragments without cloning values.
188    pub(crate) fn canonical_projection(&self) -> (Self, Vec<usize>) {
189        fn remap_slot(
190            slot: usize,
191            source_slots: &mut Vec<usize>,
192            positions: &mut HashMap<usize, usize>,
193        ) -> usize {
194            if slot == NULL_SLOT {
195                return NULL_SLOT;
196            }
197            if let Some(position) = positions.get(&slot) {
198                return *position;
199            }
200            let position = source_slots.len();
201            source_slots.push(slot);
202            positions.insert(slot, position);
203            position
204        }
205
206        let mut source_slots = Vec::new();
207        let mut positions = HashMap::new();
208        let slots = self
209            .index
210            .slots
211            .iter()
212            .map(|slot| remap_slot(*slot, &mut source_slots, &mut positions))
213            .collect();
214        let mut source_aliases = self.index.aliases.iter().collect::<Vec<_>>();
215        source_aliases.sort_unstable_by(|(left, _), (right, _)| left.cmp(right));
216        let aliases = source_aliases
217            .into_iter()
218            .map(|(name, slot)| {
219                (
220                    name.clone(),
221                    remap_slot(*slot, &mut source_slots, &mut positions),
222                )
223            })
224            .collect();
225        let mut source_internal = self.index.executor_attributes.iter().collect::<Vec<_>>();
226        source_internal.sort_unstable_by_key(|(column, _)| **column);
227        let internal = source_internal
228            .into_iter()
229            .map(|(column, slot)| {
230                (
231                    *column,
232                    remap_slot(*slot, &mut source_slots, &mut positions),
233                )
234            })
235            .collect();
236        (
237            Self::from_typed_parts_with_aliases_and_exact_precedence(
238                self.columns().to_vec(),
239                self.identities().to_vec(),
240                self.column_types().to_vec(),
241                slots,
242                source_slots.len(),
243                SchemaBuildMetadata {
244                    aliases,
245                    alias_types: self.index.cold.aliases.clone(),
246                    internal,
247                    internal_types: self.index.cold.executor_attribute_types.clone(),
248                    score_sources: self.index.cold.score_sources.clone(),
249                    wildcard_hidden: self.index.cold.wildcard_hidden.clone(),
250                    binding_only: self.index.cold.binding_only.clone(),
251                    ..SchemaBuildMetadata::default()
252                },
253            ),
254            source_slots,
255        )
256    }
257
258    /// Rebuild a schema decoded from a positional spill record.
259    ///
260    /// `None` represents a logical or alias identity whose source was absent
261    /// and therefore resolves to SQL NULL. Every physical slot is validated
262    /// before the derived lookup indexes are constructed.
263    #[expect(
264        clippy::too_many_lines,
265        reason = "projection keeps schema and physical column positions aligned"
266    )]
267    pub(crate) fn from_physical_layout(layout: PhysicalLayout) -> ExecResult<Self> {
268        let PhysicalLayout {
269            columns,
270            identities,
271            types,
272            slots,
273            physical_width,
274            aliases,
275            internal,
276            score_sources,
277            wildcard_hidden,
278        } = layout;
279        if columns.len() != slots.len() {
280            return Err(ExecError::Other(format!(
281                "physical schema has {} columns but {} logical slots",
282                columns.len(),
283                slots.len()
284            )));
285        }
286        if columns.len() != types.len() {
287            return Err(ExecError::Other(format!(
288                "physical schema has {} columns but {} logical types",
289                columns.len(),
290                types.len()
291            )));
292        }
293        if columns.len() != identities.len() {
294            return Err(ExecError::Other(format!(
295                "physical schema has {} columns but {} logical identities",
296                columns.len(),
297                identities.len()
298            )));
299        }
300        if wildcard_hidden
301            .iter()
302            .any(|position| *position >= columns.len())
303        {
304            return Err(ExecError::Other(
305                "physical schema wildcard-hidden position is outside logical width".into(),
306            ));
307        }
308        let slots = slots
309            .into_iter()
310            .map(|slot| match slot {
311                Some(slot) if slot < physical_width => Ok(slot),
312                Some(slot) => Err(ExecError::Other(format!(
313                    "physical schema logical slot {slot} is outside width {physical_width}"
314                ))),
315                None => Ok(NULL_SLOT),
316            })
317            .collect::<ExecResult<Vec<_>>>()?;
318        let mut lookup_aliases = HashMap::with_capacity(aliases.len());
319        let mut alias_types = HashMap::with_capacity(aliases.len());
320        for (identity, slot, ty) in aliases {
321            let slot = match slot {
322                Some(slot) if slot < physical_width => slot,
323                Some(slot) => {
324                    return Err(ExecError::Other(format!(
325                        "physical schema alias `{identity:?}` slot {slot} is outside width {physical_width}"
326                    )))
327                }
328                None => NULL_SLOT,
329            };
330            if lookup_aliases.insert(identity.clone(), slot).is_some() {
331                return Err(ExecError::Other(format!(
332                    "physical schema contains duplicate alias `{identity:?}`"
333                )));
334            }
335            alias_types.insert(identity, ty);
336        }
337        let mut internal_slots = HashMap::with_capacity(internal.len());
338        let mut internal_types = HashMap::with_capacity(internal.len());
339        for (column, slot, ty) in internal {
340            let slot = match slot {
341                Some(slot) if slot < physical_width => slot,
342                Some(slot) => {
343                    return Err(ExecError::Other(format!(
344                        "physical schema internal attribute `{column:?}` slot {slot} is outside width {physical_width}"
345                    )))
346                }
347                None => NULL_SLOT,
348            };
349            if internal_slots.insert(column, slot).is_some() {
350                return Err(ExecError::Other(format!(
351                    "physical schema contains duplicate internal attribute `{column:?}`"
352                )));
353            }
354            internal_types.insert(column, ty);
355        }
356        let score_sources = score_sources
357            .into_iter()
358            .map(|(qualifier, column)| {
359                if !internal_slots.contains_key(&column) {
360                    return Err(ExecError::Other(format!(
361                        "physical schema score source references missing internal attribute `{column:?}`"
362                    )));
363                }
364                Ok(ScoreSource {
365                    qualifier: qualifier.map(Box::<str>::from),
366                    column,
367                })
368            })
369            .collect::<ExecResult<Vec<_>>>()?;
370        Ok(Self::from_typed_parts_with_aliases_and_exact_precedence(
371            columns,
372            identities,
373            types,
374            slots,
375            physical_width,
376            SchemaBuildMetadata {
377                aliases: lookup_aliases,
378                alias_types,
379                internal: internal_slots,
380                internal_types,
381                score_sources,
382                wildcard_hidden,
383                binding_only: HashMap::new(),
384                ..SchemaBuildMetadata::default()
385            },
386        ))
387    }
388
389    pub(crate) fn lookup_aliases(&self) -> Vec<(&ColumnIdentity, Option<usize>)> {
390        let mut aliases = self
391            .index
392            .aliases
393            .iter()
394            .map(|(identity, slot)| (identity, (*slot != NULL_SLOT).then_some(*slot)))
395            .collect::<Vec<_>>();
396        aliases.sort_unstable_by_key(|(identity, _)| *identity);
397        aliases
398    }
399
400    pub(crate) fn lookup_aliases_with_types(
401        &self,
402    ) -> Vec<(&ColumnIdentity, Option<usize>, Option<&ColumnType>)> {
403        self.lookup_aliases()
404            .into_iter()
405            .map(|(identity, slot)| {
406                (
407                    identity,
408                    slot,
409                    self.index
410                        .cold
411                        .aliases
412                        .get(identity)
413                        .and_then(Option::as_ref),
414                )
415            })
416            .collect()
417    }
418
419    pub(crate) fn internal_columns_with_types(
420        &self,
421    ) -> Vec<(
422        uqa_sql::ast::InternalColumnRef,
423        Option<usize>,
424        Option<&ColumnType>,
425    )> {
426        let mut columns = self
427            .index
428            .executor_attributes
429            .iter()
430            .map(|(column, slot)| {
431                (
432                    *column,
433                    (*slot != NULL_SLOT).then_some(*slot),
434                    self.index
435                        .cold
436                        .executor_attribute_types
437                        .get(column)
438                        .and_then(Option::as_ref),
439                )
440            })
441            .collect::<Vec<_>>();
442        columns.sort_unstable_by_key(|(column, _, _)| *column);
443        columns
444    }
445
446    pub(crate) fn score_sources(
447        &self,
448    ) -> impl Iterator<Item = (Option<&str>, uqa_sql::ast::InternalColumnRef)> {
449        self.index
450            .cold
451            .score_sources
452            .iter()
453            .map(|source| (source.qualifier.as_deref(), source.column))
454    }
455
456    pub(crate) fn wildcard_hidden_positions(&self) -> impl Iterator<Item = usize> + '_ {
457        self.index.cold.wildcard_hidden.iter().copied()
458    }
459
460    /// Select logical input positions and attach hidden lookup identities without copying their physical values. Existing hidden aliases are retained, so nested join qualification survives another remap.
461    pub(crate) fn remap_positions(
462        input: &Self,
463        columns: &[(String, usize)],
464        aliases: &[(ColumnIdentity, usize)],
465    ) -> Self {
466        let columns = columns
467            .iter()
468            .map(|(name, logical)| (name.clone(), *logical, input.column_type(*logical).cloned()))
469            .collect::<Vec<_>>();
470        Self::remap_typed_positions(input, &columns, aliases)
471    }
472
473    /// Select logical positions with explicit output types. This is used when a binder inserts an implicit coercion and the output identity no longer has the input slot's declared type.
474    pub(crate) fn remap_typed_positions(
475        input: &Self,
476        columns: &[(String, usize, Option<ColumnType>)],
477        aliases: &[(ColumnIdentity, usize)],
478    ) -> Self {
479        let output_names = columns
480            .iter()
481            .map(|(output, _, _)| output.clone())
482            .collect::<Vec<_>>();
483        let slots = columns
484            .iter()
485            .map(|(_, logical, _)| input.slot(*logical).unwrap_or(NULL_SLOT))
486            .collect();
487        let types = columns.iter().map(|(_, _, ty)| ty.clone()).collect();
488        let identities = output_names
489            .iter()
490            .cloned()
491            .map(ColumnIdentity::unqualified)
492            .collect();
493        let wildcard_hidden = columns
494            .iter()
495            .enumerate()
496            .filter_map(|(output, (_, logical, _))| {
497                input
498                    .index
499                    .cold
500                    .wildcard_hidden
501                    .contains(logical)
502                    .then_some(output)
503            })
504            .collect();
505        let mut lookup_aliases = input.index.aliases.clone();
506        let mut alias_types = input.index.cold.aliases.clone();
507        for (identity, logical) in aliases {
508            lookup_aliases.insert(identity.clone(), input.slot(*logical).unwrap_or(NULL_SLOT));
509            alias_types.insert(identity.clone(), input.column_type(*logical).cloned());
510        }
511        Self::from_typed_parts_with_aliases_and_exact_precedence(
512            output_names,
513            identities,
514            types,
515            slots,
516            input.physical_width(),
517            SchemaBuildMetadata {
518                aliases: lookup_aliases,
519                alias_types,
520                internal: input.index.executor_attributes.clone(),
521                internal_types: input.index.cold.executor_attribute_types.clone(),
522                score_sources: input.index.cold.score_sources.clone(),
523                wildcard_hidden,
524                binding_only: input.index.cold.binding_only.clone(),
525                ..SchemaBuildMetadata::default()
526            },
527        )
528    }
529}