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    for expression in expressions {
154        collect_expr_prune_columns(expression, scope, prune, valid);
155    }
156}
157
158struct LocalTableMetadataBinding {
159    qualifier: String,
160    legacy_doc_id: bool,
161    legacy_score: bool,
162}
163
164fn single_local_table_metadata_binding(
165    catalog: &dyn AnalysisCatalog,
166    resolution: &RelationNameResolution,
167    source: &SourcePlan,
168    is_visible_cte: &dyn Fn(&str) -> bool,
169) -> Result<Option<LocalTableMetadataBinding>, SQLError> {
170    fn collect(
171        catalog: &dyn AnalysisCatalog,
172        resolution: &RelationNameResolution,
173        source: &SourcePlan,
174        is_visible_cte: &dyn Fn(&str) -> bool,
175        relations: &mut BTreeSet<(String, String)>,
176    ) -> Result<(), SQLError> {
177        match source {
178            SourcePlan::Table {
179                name,
180                qualifier,
181                alias,
182                ..
183            } => {
184                if !is_visible_cte(name) {
185                    if let Some(name) = catalog.table_name_resolved(resolution, name)? {
186                        relations.insert((alias.as_deref().unwrap_or(qualifier).to_string(), name));
187                    }
188                }
189            }
190            SourcePlan::Join { left, right, .. } => {
191                collect(catalog, resolution, left, is_visible_cte, relations)?;
192                collect(catalog, resolution, right, is_visible_cte, relations)?;
193            }
194            SourcePlan::Values { .. }
195            | SourcePlan::Function { .. }
196            | SourcePlan::FunctionGroup { .. }
197            | SourcePlan::Subquery { .. } => {}
198        }
199        Ok(())
200    }
201    if source_contains_join_alias(source) {
202        return Ok(None);
203    }
204    let mut relations = BTreeSet::new();
205    collect(catalog, resolution, source, is_visible_cte, &mut relations)?;
206    let Some((qualifier, name)) = relations.pop_first() else {
207        return Ok(None);
208    };
209    if !relations.is_empty() {
210        return Ok(None);
211    }
212    let columns = &catalog
213        .table_resolved(resolution, &name)?
214        .ok_or_else(|| SQLError::UnknownTable(name.clone()))?
215        .columns;
216    Ok(Some(LocalTableMetadataBinding {
217        qualifier,
218        legacy_doc_id: !columns.iter().any(|column| column.name == DOC_ID_COLUMN),
219        legacy_score: !columns.iter().any(|column| column.name == SCORE_COLUMN),
220    }))
221}
222
223fn source_contains_join_alias(source: &SourcePlan) -> bool {
224    match source {
225        SourcePlan::Join {
226            left, right, alias, ..
227        } => {
228            alias.is_some() || source_contains_join_alias(left) || source_contains_join_alias(right)
229        }
230        SourcePlan::Table { .. }
231        | SourcePlan::Values { .. }
232        | SourcePlan::Function { .. }
233        | SourcePlan::FunctionGroup { .. }
234        | SourcePlan::Subquery { .. } => false,
235    }
236}
237
238fn collect_join_binding_prune_columns(
239    catalog: &dyn AnalysisCatalog,
240    resolution: &RelationNameResolution,
241    from: &SourcePlan,
242    prune: &mut ColumnPrune,
243) -> Result<(), SQLError> {
244    match from {
245        SourcePlan::Join {
246            left,
247            right,
248            using,
249            natural,
250            ..
251        } => {
252            collect_join_binding_prune_columns(catalog, resolution, left, prune)?;
253            collect_join_binding_prune_columns(catalog, resolution, right, prune)?;
254            if let Some(using) = using {
255                for column in &using.columns {
256                    add_column_to_source_prune(left, column, prune);
257                    add_column_to_source_prune(right, column, prune);
258                }
259            }
260            if *natural {
261                add_all_source_columns_to_prune(catalog, resolution, left, prune)?;
262                add_all_source_columns_to_prune(catalog, resolution, right, prune)?;
263            }
264        }
265        SourcePlan::Table { .. }
266        | SourcePlan::Values { .. }
267        | SourcePlan::Function { .. }
268        | SourcePlan::FunctionGroup { .. }
269        | SourcePlan::Subquery { .. } => {}
270    }
271    Ok(())
272}
273
274fn add_column_to_source_prune(source: &SourcePlan, column: &str, prune: &mut ColumnPrune) {
275    let mut qualifiers = Vec::new();
276    collect_from_qualifiers(source, &mut qualifiers);
277    for qualifier in qualifiers {
278        if let Some(columns) = prune.get_mut(&qualifier) {
279            columns.insert(column.to_string());
280        }
281    }
282}
283
284#[expect(
285    clippy::too_many_lines,
286    reason = "preserves SELECT schema and row identity"
287)]
288fn add_all_source_columns_to_prune(
289    catalog: &dyn AnalysisCatalog,
290    resolution: &RelationNameResolution,
291    source: &SourcePlan,
292    prune: &mut ColumnPrune,
293) -> Result<(), SQLError> {
294    match source {
295        SourcePlan::Table {
296            name,
297            qualifier,
298            alias,
299            column_aliases,
300            ..
301        } => {
302            let qualifier = alias.as_deref().unwrap_or(qualifier);
303            match catalog.table_resolved(resolution, name)? {
304                Some(table) => {
305                    if let Some(columns) = prune.get_mut(qualifier) {
306                        columns.extend(table.columns.iter().enumerate().map(
307                            |(position, column)| {
308                                column_aliases
309                                    .get(position)
310                                    .cloned()
311                                    .unwrap_or_else(|| column.name.clone())
312                            },
313                        ));
314                    }
315                }
316                None => {
317                    // A CTE, view, or external relation owns its row type
318                    // outside the local table catalog. Omitting its prune
319                    // entry retains that source's complete schema.
320                    prune.remove(qualifier);
321                }
322            }
323        }
324        SourcePlan::Join { left, right, .. } => {
325            add_all_source_columns_to_prune(catalog, resolution, left, prune)?;
326            add_all_source_columns_to_prune(catalog, resolution, right, prune)?;
327        }
328        SourcePlan::Values {
329            rows,
330            alias,
331            column_aliases,
332            ..
333        } => {
334            let Some(columns) = alias.as_ref().and_then(|alias| prune.get_mut(alias)) else {
335                return Ok(());
336            };
337            if column_aliases.is_empty() {
338                columns.extend(
339                    (0..rows.first().map_or(0, Vec::len))
340                        .map(|index| format!("column{}", index + 1)),
341                );
342            } else {
343                columns.extend(column_aliases.iter().cloned());
344            }
345        }
346        SourcePlan::Function {
347            name,
348            output_name,
349            args,
350            alias,
351            column_aliases,
352            ordinality,
353            ..
354        } => {
355            let qualifier = alias.as_ref().unwrap_or(output_name);
356            let Some(columns) = prune.get_mut(qualifier) else {
357                return Ok(());
358            };
359            let routine_columns = uqa_sql::binding::catalog_sources::user_function_output_columns(
360                catalog, resolution, name,
361            )?;
362            columns.extend(routine_columns.map_or_else(
363                || {
364                    uqa_sql::semantics::table_function_empty_schema(
365                        name,
366                        output_name,
367                        alias.as_deref(),
368                        column_aliases,
369                        args.len(),
370                        *ordinality,
371                    )
372                },
373                |base| {
374                    uqa_sql::semantics::apply_table_function_aliases(
375                        base,
376                        column_aliases,
377                        *ordinality,
378                    )
379                },
380            ));
381        }
382        SourcePlan::FunctionGroup {
383            functions,
384            alias,
385            column_aliases,
386            ordinality,
387        } => {
388            let Some(qualifier) = alias
389                .as_ref()
390                .or_else(|| functions.first().map(|function| &function.output_name))
391            else {
392                return Ok(());
393            };
394            let Some(columns) = prune.get_mut(qualifier) else {
395                return Ok(());
396            };
397            let mut group_columns = Vec::new();
398            for function in functions {
399                let routine_columns =
400                    uqa_sql::binding::catalog_sources::user_function_output_columns(
401                        catalog,
402                        resolution,
403                        &function.name,
404                    )?;
405                group_columns.extend(routine_columns.map_or_else(
406                    || {
407                        uqa_sql::semantics::table_function_empty_schema(
408                            &function.name,
409                            &function.output_name,
410                            None,
411                            &function.column_aliases,
412                            function.args.len(),
413                            false,
414                        )
415                    },
416                    |base| {
417                        uqa_sql::semantics::apply_table_function_aliases(
418                            base,
419                            &function.column_aliases,
420                            false,
421                        )
422                    },
423                ));
424            }
425            if *ordinality {
426                group_columns.push("ordinality".into());
427            }
428            for (column, alias) in group_columns.iter_mut().zip(column_aliases) {
429                column.clone_from(alias);
430            }
431            columns.extend(group_columns);
432        }
433        SourcePlan::Subquery {
434            body,
435            alias,
436            column_aliases,
437        } => {
438            let Some(columns) = alias.as_ref().and_then(|alias| prune.get_mut(alias)) else {
439                return Ok(());
440            };
441            if column_aliases.is_empty() {
442                columns.extend(
443                    uqa_sql::semantics::query_plan_output_columns(body).unwrap_or_default(),
444                );
445            } else {
446                columns.extend(column_aliases.iter().cloned());
447            }
448        }
449    }
450    Ok(())
451}
452
453pub use uqa_sql::semantics::collect_from_qualifiers;
454
455fn collect_from_prune_columns(
456    from: &SourcePlan,
457    scope: PruneScope<'_>,
458    prune: &mut ColumnPrune,
459    valid: &mut bool,
460) {
461    match from {
462        SourcePlan::Join {
463            left, right, on, ..
464        } => {
465            collect_from_prune_columns(left, scope, prune, valid);
466            collect_from_prune_columns(right, scope, prune, valid);
467            if let Some(on) = on.as_ref() {
468                collect_expr_prune_columns(on, scope, prune, valid);
469            }
470        }
471        SourcePlan::Values { rows, .. } => {
472            for row in rows {
473                for expr in row {
474                    collect_expr_prune_columns(expr, scope, prune, valid);
475                }
476            }
477        }
478        SourcePlan::Function { args, .. } => {
479            for expr in args {
480                collect_expr_prune_columns(expr, scope, prune, valid);
481            }
482        }
483        SourcePlan::FunctionGroup { functions, .. } => {
484            for function in functions {
485                for expr in &function.args {
486                    collect_expr_prune_columns(expr, scope, prune, valid);
487                }
488            }
489        }
490        SourcePlan::Subquery { .. } => {
491            *valid = false;
492        }
493        SourcePlan::Table { .. } => {}
494    }
495}
496
497#[expect(
498    clippy::too_many_lines,
499    reason = "preserves SELECT schema and row identity"
500)]
501fn collect_expr_prune_columns(
502    expr: &ScalarExpr,
503    scope: PruneScope<'_>,
504    prune: &mut ColumnPrune,
505    valid: &mut bool,
506) {
507    match expr {
508        ScalarExpr::Column(column) => {
509            for qualifier in scope.qualifiers {
510                if qualifier.eq_ignore_ascii_case(column) {
511                    let Some(source) = prune.get_mut(qualifier) else {
512                        *valid = false;
513                        return;
514                    };
515                    source.retain_all();
516                }
517            }
518            if let Some(qualifier) = scope.metadata_qualifier {
519                let metadata = match column.as_str() {
520                    DOC_ID_COLUMN if scope.legacy_doc_id => Some(true),
521                    SCORE_COLUMN if scope.legacy_score => Some(false),
522                    _ => None,
523                };
524                if let Some(doc_id) = metadata {
525                    let Some(source) = prune.get_mut(qualifier) else {
526                        *valid = false;
527                        return;
528                    };
529                    source.insert(column.clone());
530                    if doc_id {
531                        source.metadata_mut().request_doc_id();
532                    } else {
533                        source.metadata_mut().request_score();
534                    }
535                    return;
536                }
537            }
538            for qualifier in scope.qualifiers {
539                if let Some(columns) = prune.get_mut(qualifier) {
540                    columns.insert(column.clone());
541                }
542            }
543        }
544        ScalarExpr::QualifiedColumn {
545            qualifier, column, ..
546        } => {
547            if scope.metadata_qualifier == Some(qualifier.as_str()) {
548                let metadata = match column.as_str() {
549                    DOC_ID_COLUMN if scope.legacy_doc_id => Some(true),
550                    SCORE_COLUMN if scope.legacy_score => Some(false),
551                    _ => None,
552                };
553                if let Some(doc_id) = metadata {
554                    let Some(source) = prune.get_mut(qualifier) else {
555                        *valid = false;
556                        return;
557                    };
558                    source.insert(column.clone());
559                    if doc_id {
560                        source.metadata_mut().request_doc_id();
561                    } else {
562                        source.metadata_mut().request_score();
563                    }
564                    return;
565                }
566            }
567            if qualifier == META_QUALIFIER && !prune.contains_key(META_QUALIFIER) {
568                let Some(source) = scope
569                    .metadata_qualifier
570                    .and_then(|source| prune.get_mut(source))
571                else {
572                    *valid = false;
573                    return;
574                };
575                match column.as_str() {
576                    META_DOC_ID_COLUMN => source.metadata_mut().request_doc_id(),
577                    META_SCORE_COLUMN => source.metadata_mut().request_score(),
578                    _ => *valid = false,
579                }
580                return;
581            }
582            if let Some(columns) = prune.get_mut(qualifier) {
583                columns.insert(column.clone());
584            } else {
585                *valid = false;
586            }
587        }
588        ScalarExpr::Literal(_) | ScalarExpr::TypedLiteral { .. } | ScalarExpr::Param(_) => {}
589        ScalarExpr::Default
590        | ScalarExpr::Star
591        | ScalarExpr::Position(_)
592        | ScalarExpr::InternalColumn(_)
593        | ScalarExpr::QualifiedStar(_)
594        | ScalarExpr::ScalarSubquery(_)
595        | ScalarExpr::Exists { .. } => {
596            *valid = false;
597        }
598        ScalarExpr::Array(items)
599        | ScalarExpr::Row(items)
600        | ScalarExpr::And(items)
601        | ScalarExpr::Or(items) => {
602            for item in items {
603                collect_expr_prune_columns(item, scope, prune, valid);
604            }
605        }
606        ScalarExpr::Func {
607            name,
608            args,
609            order_by,
610            filter,
611            ..
612        } => {
613            // 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.
614            if !(name.eq_ignore_ascii_case("count")
615                && matches!(args.as_slice(), [ScalarExpr::Star]))
616            {
617                for arg in args {
618                    collect_expr_prune_columns(arg, scope, prune, valid);
619                }
620            }
621            for order in order_by {
622                collect_expr_prune_columns(&order.expr, scope, prune, valid);
623            }
624            if let Some(filter) = filter.as_ref() {
625                collect_expr_prune_columns(filter, scope, prune, valid);
626            }
627        }
628        ScalarExpr::Binary { lhs, rhs, .. } => {
629            collect_expr_prune_columns(lhs, scope, prune, valid);
630            collect_expr_prune_columns(rhs, scope, prune, valid);
631        }
632        ScalarExpr::Not(inner)
633        | ScalarExpr::UnaryMinus(inner)
634        | ScalarExpr::IsNull { expr: inner, .. }
635        | ScalarExpr::Cast { expr: inner, .. } => {
636            collect_expr_prune_columns(inner, scope, prune, valid);
637        }
638        ScalarExpr::Between { expr, low, high } => {
639            collect_expr_prune_columns(expr, scope, prune, valid);
640            collect_expr_prune_columns(low, scope, prune, valid);
641            collect_expr_prune_columns(high, scope, prune, valid);
642        }
643        ScalarExpr::InList { expr, list, .. } => {
644            collect_expr_prune_columns(expr, scope, prune, valid);
645            for item in list {
646                collect_expr_prune_columns(item, scope, prune, valid);
647            }
648        }
649        ScalarExpr::WindowCall { args, spec, .. } => {
650            for argument in args {
651                collect_expr_prune_columns(argument, scope, prune, valid);
652            }
653            for expression in &spec.partition_by {
654                collect_expr_prune_columns(expression, scope, prune, valid);
655            }
656            for order in &spec.order_by {
657                collect_expr_prune_columns(&order.expr, scope, prune, valid);
658            }
659            if let Some(frame) = &spec.frame {
660                for bound in [&frame.start, &frame.end] {
661                    if let uqa_sql::ScalarFrameBound::Preceding(expression)
662                    | uqa_sql::ScalarFrameBound::Following(expression) = bound
663                    {
664                        collect_expr_prune_columns(expression, scope, prune, valid);
665                    }
666                }
667            }
668            *valid = false;
669        }
670        ScalarExpr::InSubquery { expr, .. } => {
671            collect_expr_prune_columns(expr, scope, prune, valid);
672            *valid = false;
673        }
674        ScalarExpr::Case {
675            base,
676            when,
677            else_branch,
678        } => {
679            if let Some(base) = base.as_ref() {
680                collect_expr_prune_columns(base, scope, prune, valid);
681            }
682            for (cond, result) in when {
683                collect_expr_prune_columns(cond, scope, prune, valid);
684                collect_expr_prune_columns(result, scope, prune, valid);
685            }
686            if let Some(else_branch) = else_branch.as_ref() {
687                collect_expr_prune_columns(else_branch, scope, prune, valid);
688            }
689        }
690    }
691}
uqa_planner/column_pruning.rs

uqa_planner/
column_pruning.rs
