uqa_planner/
column_pruning.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Required source columns and row metadata for physical query projection.
8
9use std::collections::BTreeSet;
10use uqa_sql::{
11    catalog::{analysis::AnalysisCatalog, resolution::RelationNameResolution},
12    plan::{
13        source_projection::{ColumnPrune, SourceProjection},
14        QueryBlockPlan, SourcePlan,
15    },
16    semantics::volatility::{expr_contains_volatile_function, VolatilityCatalog},
17    semantics::{
18        expr_contains_subquery, DOC_ID_COLUMN, META_DOC_ID_COLUMN, META_QUALIFIER,
19        META_SCORE_COLUMN, SCORE_COLUMN,
20    },
21    SQLError, ScalarExpr,
22};
23
24#[cfg(test)]
25mod tests;
26
27#[derive(Clone, Copy)]
28pub struct ColumnPruneContext<'a> {
29    pub catalog: &'a dyn AnalysisCatalog,
30    pub resolution: &'a RelationNameResolution,
31    pub volatility: &'a dyn VolatilityCatalog,
32    pub is_visible_cte: &'a dyn Fn(&str) -> bool,
33}
34
35fn has_window(projections: &[uqa_sql::plan::ProjectionPlan]) -> bool {
36    projections
37        .iter()
38        .any(|projection| uqa_sql::semantics::windows::expr_has_window(&projection.expr))
39}
40
41pub fn column_prune_for_stmt(
42    context: ColumnPruneContext<'_>,
43    stmt: &QueryBlockPlan,
44    from: &SourcePlan,
45) -> Result<Option<ColumnPrune>, SQLError> {
46    column_prune_for_stmt_with_filter(context, stmt, from, stmt.r#where.as_ref())
47}
48
49/// Compute the document projection for `stmt` while treating `filter` as the
50/// only predicate that remains to be evaluated by the relational pipeline.
51/// Accelerated retrieval consumes its search predicate before constructing a
52/// scored document source, so its field
53/// arguments are index dependencies rather than row-materialization
54/// dependencies. Callers that have executed retrieval pass only the residual
55/// predicate here; ordinary scans retain the statement's original `WHERE` via
56/// [`column_prune_for_stmt`].
57pub fn column_prune_for_stmt_with_filter(
58    context: ColumnPruneContext<'_>,
59    stmt: &QueryBlockPlan,
60    from: &SourcePlan,
61    filter: Option<&ScalarExpr>,
62) -> Result<Option<ColumnPrune>, SQLError> {
63    let catalog = context.catalog;
64    let resolution = context.resolution;
65    let requires_full_projection = source_contains_join_alias(from)
66        || has_window(&stmt.projections)
67        || stmt.projections.iter().any(|projection| {
68            matches!(projection.expr, ScalarExpr::Star)
69                || expr_contains_subquery(&projection.expr)
70                || expr_contains_volatile_function(context.volatility, &projection.expr)
71        });
72
73    let mut qualifiers = Vec::new();
74    collect_from_qualifiers(from, &mut qualifiers);
75    if qualifiers.is_empty() {
76        return Ok(None);
77    }
78
79    let metadata_binding =
80        single_local_table_metadata_binding(catalog, resolution, from, context.is_visible_cte)?;
81    let scope = PruneScope {
82        qualifiers: &qualifiers,
83        metadata_qualifier: metadata_binding
84            .as_ref()
85            .map(|binding| binding.qualifier.as_str()),
86        legacy_doc_id: metadata_binding
87            .as_ref()
88            .is_some_and(|binding| binding.legacy_doc_id),
89        legacy_score: metadata_binding
90            .as_ref()
91            .is_some_and(|binding| binding.legacy_score),
92    };
93    let mut prune: ColumnPrune = qualifiers
94        .iter()
95        .map(|qualifier| {
96            (
97                qualifier.clone(),
98                if requires_full_projection {
99                    SourceProjection::retaining_all()
100                } else {
101                    SourceProjection::default()
102                },
103            )
104        })
105        .collect();
106    let mut valid = true;
107    collect_from_prune_columns(from, scope, &mut prune, &mut valid);
108    collect_join_binding_prune_columns(catalog, resolution, from, &mut prune)?;
109    collect_query_block_prune_columns(stmt, filter, scope, &mut prune, &mut valid);
110    let metadata_requested = prune
111        .values()
112        .any(|projection| !projection.metadata().is_empty());
113    if requires_full_projection {
114        return Ok(metadata_requested.then_some(prune));
115    }
116    if !valid {
117        if metadata_requested {
118            for projection in prune.values_mut() {
119                projection.retain_all();
120            }
121            return Ok(Some(prune));
122        }
123        return Ok(None);
124    }
125    Ok(Some(prune))
126}
127
128#[derive(Clone, Copy)]
129struct PruneScope<'a> {
130    qualifiers: &'a [String],
131    metadata_qualifier: Option<&'a str>,
132    legacy_doc_id: bool,
133    legacy_score: bool,
134}
135
136fn collect_query_block_prune_columns(
137    stmt: &QueryBlockPlan,
138    filter: Option<&ScalarExpr>,
139    scope: PruneScope<'_>,
140    prune: &mut ColumnPrune,
141    valid: &mut bool,
142) {
143    let expressions = stmt
144        .projections
145        .iter()
146        .map(|projection| &projection.expr)
147        .chain(filter)
148        .chain(stmt.group_by.iter())
149        .chain(stmt.grouping_sets.iter().flatten())
150        .chain(stmt.having.iter())
151        .chain(stmt.order_by.iter().map(|order| &order.expr))
152        .chain(stmt.distinct_on.iter())
153        .chain(
154            stmt.windows
155                .iter()
156                .flat_map(|window| window.spec.expressions()),
157        );
158    for expression in expressions {
159        collect_expr_prune_columns(expression, scope, prune, valid);
160    }
161}
162
163struct LocalTableMetadataBinding {
164    qualifier: String,
165    legacy_doc_id: bool,
166    legacy_score: bool,
167}
168
169fn single_local_table_metadata_binding(
170    catalog: &dyn AnalysisCatalog,
171    resolution: &RelationNameResolution,
172    source: &SourcePlan,
173    is_visible_cte: &dyn Fn(&str) -> bool,
174) -> Result<Option<LocalTableMetadataBinding>, SQLError> {
175    fn collect(
176        catalog: &dyn AnalysisCatalog,
177        resolution: &RelationNameResolution,
178        source: &SourcePlan,
179        is_visible_cte: &dyn Fn(&str) -> bool,
180        relations: &mut BTreeSet<(String, String)>,
181    ) -> Result<(), SQLError> {
182        match source {
183            SourcePlan::Table {
184                name,
185                qualifier,
186                alias,
187                ..
188            } => {
189                if !is_visible_cte(name) {
190                    if let Some(name) = catalog.table_name_resolved(resolution, name)? {
191                        relations.insert((alias.as_deref().unwrap_or(qualifier).to_string(), name));
192                    }
193                }
194            }
195            SourcePlan::Join { left, right, .. } => {
196                collect(catalog, resolution, left, is_visible_cte, relations)?;
197                collect(catalog, resolution, right, is_visible_cte, relations)?;
198            }
199            SourcePlan::Values { .. }
200            | SourcePlan::Function { .. }
201            | SourcePlan::FunctionGroup { .. }
202            | SourcePlan::Subquery { .. } => {}
203        }
204        Ok(())
205    }
206    if source_contains_join_alias(source) {
207        return Ok(None);
208    }
209    let mut relations = BTreeSet::new();
210    collect(catalog, resolution, source, is_visible_cte, &mut relations)?;
211    let Some((qualifier, name)) = relations.pop_first() else {
212        return Ok(None);
213    };
214    if !relations.is_empty() {
215        return Ok(None);
216    }
217    let columns = &catalog
218        .table_resolved(resolution, &name)?
219        .ok_or_else(|| SQLError::UnknownTable(name.clone()))?
220        .columns;
221    Ok(Some(LocalTableMetadataBinding {
222        qualifier,
223        legacy_doc_id: !columns.iter().any(|column| column.name == DOC_ID_COLUMN),
224        legacy_score: !columns.iter().any(|column| column.name == SCORE_COLUMN),
225    }))
226}
227
228fn source_contains_join_alias(source: &SourcePlan) -> bool {
229    match source {
230        SourcePlan::Join {
231            left, right, alias, ..
232        } => {
233            alias.is_some() || source_contains_join_alias(left) || source_contains_join_alias(right)
234        }
235        SourcePlan::Table { .. }
236        | SourcePlan::Values { .. }
237        | SourcePlan::Function { .. }
238        | SourcePlan::FunctionGroup { .. }
239        | SourcePlan::Subquery { .. } => false,
240    }
241}
242
243fn collect_join_binding_prune_columns(
244    catalog: &dyn AnalysisCatalog,
245    resolution: &RelationNameResolution,
246    from: &SourcePlan,
247    prune: &mut ColumnPrune,
248) -> Result<(), SQLError> {
249    match from {
250        SourcePlan::Join {
251            left,
252            right,
253            using,
254            natural,
255            ..
256        } => {
257            collect_join_binding_prune_columns(catalog, resolution, left, prune)?;
258            collect_join_binding_prune_columns(catalog, resolution, right, prune)?;
259            if let Some(using) = using {
260                for column in &using.columns {
261                    add_column_to_source_prune(left, column, prune);
262                    add_column_to_source_prune(right, column, prune);
263                }
264            }
265            if *natural {
266                add_all_source_columns_to_prune(catalog, resolution, left, prune)?;
267                add_all_source_columns_to_prune(catalog, resolution, right, prune)?;
268            }
269        }
270        SourcePlan::Table { .. }
271        | SourcePlan::Values { .. }
272        | SourcePlan::Function { .. }
273        | SourcePlan::FunctionGroup { .. }
274        | SourcePlan::Subquery { .. } => {}
275    }
276    Ok(())
277}
278
279fn add_column_to_source_prune(source: &SourcePlan, column: &str, prune: &mut ColumnPrune) {
280    let mut qualifiers = Vec::new();
281    collect_from_qualifiers(source, &mut qualifiers);
282    for qualifier in qualifiers {
283        if let Some(columns) = prune.get_mut(&qualifier) {
284            columns.insert(column.to_string());
285        }
286    }
287}
288
289#[expect(
290    clippy::too_many_lines,
291    reason = "preserves SELECT schema and row identity"
292)]
293fn add_all_source_columns_to_prune(
294    catalog: &dyn AnalysisCatalog,
295    resolution: &RelationNameResolution,
296    source: &SourcePlan,
297    prune: &mut ColumnPrune,
298) -> Result<(), SQLError> {
299    match source {
300        SourcePlan::Table {
301            name,
302            qualifier,
303            alias,
304            column_aliases,
305            ..
306        } => {
307            let qualifier = alias.as_deref().unwrap_or(qualifier);
308            match catalog.table_resolved(resolution, name)? {
309                Some(table) => {
310                    if let Some(columns) = prune.get_mut(qualifier) {
311                        columns.extend(table.columns.iter().enumerate().map(
312                            |(position, column)| {
313                                column_aliases
314                                    .get(position)
315                                    .cloned()
316                                    .unwrap_or_else(|| column.name.clone())
317                            },
318                        ));
319                    }
320                }
321                None => {
322                    // A CTE, view, or external relation owns its row type
323                    // outside the local table catalog. Omitting its prune
324                    // entry retains that source's complete schema.
325                    prune.remove(qualifier);
326                }
327            }
328        }
329        SourcePlan::Join { left, right, .. } => {
330            add_all_source_columns_to_prune(catalog, resolution, left, prune)?;
331            add_all_source_columns_to_prune(catalog, resolution, right, prune)?;
332        }
333        SourcePlan::Values {
334            rows,
335            alias,
336            column_aliases,
337            ..
338        } => {
339            let Some(columns) = alias.as_ref().and_then(|alias| prune.get_mut(alias)) else {
340                return Ok(());
341            };
342            if column_aliases.is_empty() {
343                columns.extend(
344                    (0..rows.first().map_or(0, Vec::len))
345                        .map(|index| format!("column{}", index + 1)),
346                );
347            } else {
348                columns.extend(column_aliases.iter().cloned());
349            }
350        }
351        SourcePlan::Function {
352            name,
353            output_name,
354            args,
355            alias,
356            column_aliases,
357            ordinality,
358            ..
359        } => {
360            let qualifier = alias.as_ref().unwrap_or(output_name);
361            let Some(columns) = prune.get_mut(qualifier) else {
362                return Ok(());
363            };
364            let routine_columns = uqa_sql::binding::catalog_sources::user_function_output_columns(
365                catalog, resolution, name,
366            )?;
367            columns.extend(routine_columns.map_or_else(
368                || {
369                    uqa_sql::semantics::table_function_empty_schema(
370                        name,
371                        output_name,
372                        alias.as_deref(),
373                        column_aliases,
374                        args.len(),
375                        *ordinality,
376                    )
377                },
378                |base| {
379                    uqa_sql::semantics::apply_table_function_aliases(
380                        base,
381                        column_aliases,
382                        *ordinality,
383                    )
384                },
385            ));
386        }
387        SourcePlan::FunctionGroup {
388            functions,
389            alias,
390            column_aliases,
391            ordinality,
392        } => {
393            let Some(qualifier) = alias
394                .as_ref()
395                .or_else(|| functions.first().map(|function| &function.output_name))
396            else {
397                return Ok(());
398            };
399            let Some(columns) = prune.get_mut(qualifier) else {
400                return Ok(());
401            };
402            let mut group_columns = Vec::new();
403            for function in functions {
404                let routine_columns =
405                    uqa_sql::binding::catalog_sources::user_function_output_columns(
406                        catalog,
407                        resolution,
408                        &function.name,
409                    )?;
410                group_columns.extend(routine_columns.map_or_else(
411                    || {
412                        uqa_sql::semantics::table_function_empty_schema(
413                            &function.name,
414                            &function.output_name,
415                            None,
416                            &function.column_aliases,
417                            function.args.len(),
418                            false,
419                        )
420                    },
421                    |base| {
422                        uqa_sql::semantics::apply_table_function_aliases(
423                            base,
424                            &function.column_aliases,
425                            false,
426                        )
427                    },
428                ));
429            }
430            if *ordinality {
431                group_columns.push("ordinality".into());
432            }
433            for (column, alias) in group_columns.iter_mut().zip(column_aliases) {
434                column.clone_from(alias);
435            }
436            columns.extend(group_columns);
437        }
438        SourcePlan::Subquery {
439            body,
440            alias,
441            column_aliases,
442        } => {
443            let Some(columns) = alias.as_ref().and_then(|alias| prune.get_mut(alias)) else {
444                return Ok(());
445            };
446            if column_aliases.is_empty() {
447                columns.extend(
448                    uqa_sql::semantics::query_plan_output_columns(body).unwrap_or_default(),
449                );
450            } else {
451                columns.extend(column_aliases.iter().cloned());
452            }
453        }
454    }
455    Ok(())
456}
457
458pub use uqa_sql::semantics::collect_from_qualifiers;
459
460fn collect_from_prune_columns(
461    from: &SourcePlan,
462    scope: PruneScope<'_>,
463    prune: &mut ColumnPrune,
464    valid: &mut bool,
465) {
466    match from {
467        SourcePlan::Join {
468            left, right, on, ..
469        } => {
470            collect_from_prune_columns(left, scope, prune, valid);
471            collect_from_prune_columns(right, scope, prune, valid);
472            if let Some(on) = on.as_ref() {
473                collect_expr_prune_columns(on, scope, prune, valid);
474            }
475        }
476        SourcePlan::Values { rows, .. } => {
477            for row in rows {
478                for expr in row {
479                    collect_expr_prune_columns(expr, scope, prune, valid);
480                }
481            }
482        }
483        SourcePlan::Function { args, .. } => {
484            for expr in args {
485                collect_expr_prune_columns(expr, scope, prune, valid);
486            }
487        }
488        SourcePlan::FunctionGroup { functions, .. } => {
489            for function in functions {
490                for expr in &function.args {
491                    collect_expr_prune_columns(expr, scope, prune, valid);
492                }
493            }
494        }
495        SourcePlan::Subquery { .. } => {
496            *valid = false;
497        }
498        SourcePlan::Table { .. } => {}
499    }
500}
501
502#[expect(
503    clippy::too_many_lines,
504    reason = "preserves SELECT schema and row identity"
505)]
506fn collect_expr_prune_columns(
507    expr: &ScalarExpr,
508    scope: PruneScope<'_>,
509    prune: &mut ColumnPrune,
510    valid: &mut bool,
511) {
512    match expr {
513        ScalarExpr::Column(column) => {
514            for qualifier in scope.qualifiers {
515                if qualifier.eq_ignore_ascii_case(column) {
516                    let Some(source) = prune.get_mut(qualifier) else {
517                        *valid = false;
518                        return;
519                    };
520                    source.retain_all();
521                }
522            }
523            if let Some(qualifier) = scope.metadata_qualifier {
524                let metadata = match column.as_str() {
525                    DOC_ID_COLUMN if scope.legacy_doc_id => Some(true),
526                    SCORE_COLUMN if scope.legacy_score => Some(false),
527                    _ => None,
528                };
529                if let Some(doc_id) = metadata {
530                    let Some(source) = prune.get_mut(qualifier) else {
531                        *valid = false;
532                        return;
533                    };
534                    source.insert(column.clone());
535                    if doc_id {
536                        source.metadata_mut().request_doc_id();
537                    } else {
538                        source.metadata_mut().request_score();
539                    }
540                    return;
541                }
542            }
543            for qualifier in scope.qualifiers {
544                if let Some(columns) = prune.get_mut(qualifier) {
545                    columns.insert(column.clone());
546                }
547            }
548        }
549        ScalarExpr::QualifiedColumn {
550            qualifier, column, ..
551        } => {
552            if scope.metadata_qualifier == Some(qualifier.as_str()) {
553                let metadata = match column.as_str() {
554                    DOC_ID_COLUMN if scope.legacy_doc_id => Some(true),
555                    SCORE_COLUMN if scope.legacy_score => Some(false),
556                    _ => None,
557                };
558                if let Some(doc_id) = metadata {
559                    let Some(source) = prune.get_mut(qualifier) else {
560                        *valid = false;
561                        return;
562                    };
563                    source.insert(column.clone());
564                    if doc_id {
565                        source.metadata_mut().request_doc_id();
566                    } else {
567                        source.metadata_mut().request_score();
568                    }
569                    return;
570                }
571            }
572            if qualifier == META_QUALIFIER && !prune.contains_key(META_QUALIFIER) {
573                let Some(source) = scope
574                    .metadata_qualifier
575                    .and_then(|source| prune.get_mut(source))
576                else {
577                    *valid = false;
578                    return;
579                };
580                match column.as_str() {
581                    META_DOC_ID_COLUMN => source.metadata_mut().request_doc_id(),
582                    META_SCORE_COLUMN => source.metadata_mut().request_score(),
583                    _ => *valid = false,
584                }
585                return;
586            }
587            if let Some(columns) = prune.get_mut(qualifier) {
588                columns.insert(column.clone());
589            } else {
590                *valid = false;
591            }
592        }
593        ScalarExpr::Literal(_) | ScalarExpr::TypedLiteral { .. } | ScalarExpr::Param(_) => {}
594        ScalarExpr::Default
595        | ScalarExpr::Star
596        | ScalarExpr::Position(_)
597        | ScalarExpr::InternalColumn(_)
598        | ScalarExpr::QualifiedStar(_)
599        | ScalarExpr::ScalarSubquery(_)
600        | ScalarExpr::Exists { .. } => {
601            *valid = false;
602        }
603        ScalarExpr::Array(items)
604        | ScalarExpr::Row(items)
605        | ScalarExpr::CompositeRow { items, .. }
606        | ScalarExpr::And(items)
607        | ScalarExpr::Or(items) => {
608            for item in items {
609                collect_expr_prune_columns(item, scope, prune, valid);
610            }
611        }
612        ScalarExpr::Func {
613            name,
614            args,
615            order_by,
616            filter,
617            ..
618        } => {
619            // COUNT(*) consumes one unit per input row; the star is not a whole-row value. Its FILTER and aggregate ordering still depend on ordinary source columns.
620            if !(name.eq_ignore_ascii_case("count")
621                && matches!(args.as_slice(), [ScalarExpr::Star]))
622            {
623                for arg in args {
624                    collect_expr_prune_columns(arg, scope, prune, valid);
625                }
626            }
627            for order in order_by {
628                collect_expr_prune_columns(&order.expr, scope, prune, valid);
629            }
630            if let Some(filter) = filter.as_ref() {
631                collect_expr_prune_columns(filter, scope, prune, valid);
632            }
633        }
634        ScalarExpr::Binary { lhs, rhs, .. } => {
635            collect_expr_prune_columns(lhs, scope, prune, valid);
636            collect_expr_prune_columns(rhs, scope, prune, valid);
637        }
638        ScalarExpr::Not(inner)
639        | ScalarExpr::UnaryMinus(inner)
640        | ScalarExpr::IsNull { expr: inner, .. }
641        | ScalarExpr::Cast { expr: inner, .. } => {
642            collect_expr_prune_columns(inner, scope, prune, valid);
643        }
644        ScalarExpr::Between { expr, low, high } => {
645            collect_expr_prune_columns(expr, scope, prune, valid);
646            collect_expr_prune_columns(low, scope, prune, valid);
647            collect_expr_prune_columns(high, scope, prune, valid);
648        }
649        ScalarExpr::InList { expr, list, .. } => {
650            collect_expr_prune_columns(expr, scope, prune, valid);
651            for item in list {
652                collect_expr_prune_columns(item, scope, prune, valid);
653            }
654        }
655        ScalarExpr::WindowCall {
656            args, spec, filter, ..
657        } => {
658            for argument in args {
659                collect_expr_prune_columns(argument, scope, prune, valid);
660            }
661            if let Some(filter) = filter {
662                collect_expr_prune_columns(filter, scope, prune, valid);
663            }
664            for expression in &spec.partition_by {
665                collect_expr_prune_columns(expression, scope, prune, valid);
666            }
667            for order in &spec.order_by {
668                collect_expr_prune_columns(&order.expr, scope, prune, valid);
669            }
670            if let Some(frame) = &spec.frame {
671                for bound in [&frame.start, &frame.end] {
672                    if let uqa_sql::ScalarFrameBound::Preceding(expression)
673                    | uqa_sql::ScalarFrameBound::Following(expression) = bound
674                    {
675                        collect_expr_prune_columns(expression, scope, prune, valid);
676                    }
677                }
678            }
679            *valid = false;
680        }
681        ScalarExpr::InSubquery { expr, .. } => {
682            collect_expr_prune_columns(expr, scope, prune, valid);
683            *valid = false;
684        }
685        ScalarExpr::Case {
686            base,
687            when,
688            else_branch,
689        } => {
690            if let Some(base) = base.as_ref() {
691                collect_expr_prune_columns(base, scope, prune, valid);
692            }
693            for (cond, result) in when {
694                collect_expr_prune_columns(cond, scope, prune, valid);
695                collect_expr_prune_columns(result, scope, prune, valid);
696            }
697            if let Some(else_branch) = else_branch.as_ref() {
698                collect_expr_prune_columns(else_branch, scope, prune, valid);
699            }
700        }
701    }
702}
uqa_planner/column_pruning.rs

uqa_planner/
column_pruning.rs
