sqlglot.optimizer.pushdown_projections

  1from __future__ import annotations
  2
  3import typing as t
  4from collections import defaultdict
  5
  6from sqlglot import alias, exp
  7from sqlglot.optimizer.qualify_columns import Resolver
  8from sqlglot.optimizer.scope import Scope, traverse_scope
  9from sqlglot.schema import ensure_schema
 10from sqlglot.errors import OptimizeError
 11from sqlglot.helper import seq_get
 12
 13if t.TYPE_CHECKING:
 14    from sqlglot._typing import E
 15    from sqlglot.schema import Schema
 16    from sqlglot.dialects.dialect import DialectType
 17
 18# Sentinel value that means an outer query selecting ALL columns
 19SELECT_ALL = object()
 20
 21
 22# Selection to use if selection list is empty
 23def default_selection(is_agg: bool) -> exp.Alias:
 24    return alias(exp.Max(this=exp.Literal.number(1)) if is_agg else "1", "_").assert_is(exp.Alias)
 25
 26
 27def pushdown_projections(
 28    expression: E,
 29    schema: dict[str, object] | Schema | None = None,
 30    remove_unused_selections: bool = True,
 31    dialect: DialectType = None,
 32) -> E:
 33    """
 34    Rewrite sqlglot AST to remove unused columns projections.
 35
 36    Example:
 37        >>> import sqlglot
 38        >>> sql = "SELECT y.a AS a FROM (SELECT x.a AS a, x.b AS b FROM x) AS y"
 39        >>> expression = sqlglot.parse_one(sql)
 40        >>> pushdown_projections(expression).sql()
 41        'SELECT y.a AS a FROM (SELECT x.a AS a FROM x) AS y'
 42
 43    Args:
 44        expression (sqlglot.Expr): expression to optimize
 45        remove_unused_selections (bool): remove selects that are unused
 46    Returns:
 47        sqlglot.Expr: optimized expression
 48    """
 49    # Map of Scope to all columns being selected by outer queries.
 50    schema = ensure_schema(schema, dialect=dialect)
 51    source_column_alias_count: dict[exp.Expr | Scope, int] = {}
 52    referenced_columns: defaultdict[Scope, set[str | object]] = defaultdict(set)
 53
 54    # We build the scope tree (which is traversed in DFS postorder), then iterate
 55    # over the result in reverse order. This should ensure that the set of selected
 56    # columns for a particular scope are completely build by the time we get to it.
 57    for scope in reversed(traverse_scope(expression)):
 58        parent_selections = referenced_columns.get(scope, {SELECT_ALL})
 59        alias_count = source_column_alias_count.get(scope, 0)
 60
 61        # We can't remove columns SELECT DISTINCT nor UNION DISTINCT.
 62        if scope.expression.args.get("distinct"):
 63            parent_selections = {SELECT_ALL}
 64
 65        if isinstance(scope.expression, exp.SetOperation):
 66            set_op = scope.expression
 67            if set_op.kind or set_op.side:
 68                # Do not optimize this set operation if it's using the BigQuery specific
 69                # kind / side syntax (e.g INNER UNION ALL BY NAME) which changes the semantics of the operation
 70                continue
 71
 72            left, right = scope.union_scopes
 73            le = left.expression
 74            re = right.expression
 75
 76            if not (isinstance(le, exp.Selectable) and isinstance(re, exp.Selectable)):
 77                continue
 78
 79            if len(le.selects) != len(re.selects):
 80                scope_sql = scope.expression.sql(dialect=dialect)
 81                raise OptimizeError(f"Invalid set operation due to column mismatch: {scope_sql}.")
 82
 83            referenced_columns[left] = parent_selections
 84
 85            if re.is_star:
 86                referenced_columns[right] = parent_selections
 87            elif not le.is_star:
 88                if scope.expression.args.get("by_name"):
 89                    referenced_columns[right] = referenced_columns[left]
 90                else:
 91                    referenced_columns[right] = {
 92                        re.selects[i].alias_or_name
 93                        for i, select in enumerate(le.selects)
 94                        if SELECT_ALL in parent_selections
 95                        or select.alias_or_name in parent_selections
 96                    }
 97
 98        if isinstance(scope.expression, exp.Select):
 99            if remove_unused_selections:
100                _remove_unused_selections(scope, parent_selections, schema, alias_count)
101
102            if scope.expression.is_star:
103                continue
104
105            # Group columns by source name
106            selects: dict[str, set[object]] = defaultdict(set)
107            for col in scope.columns:
108                table_name = col.table
109                col_name = col.name
110                selects[table_name].add(col_name)
111
112            # Push the selected columns down to the next scope
113            for name, (node, source) in scope.selected_sources.items():
114                if isinstance(source, Scope) and isinstance(source.expression, exp.Selectable):
115                    select = seq_get(source.expression.selects, 0)
116
117                    if scope.pivots or isinstance(select, exp.QueryTransform):
118                        columns: set[object] = {SELECT_ALL}
119                    else:
120                        columns = selects.get(name) or set()
121
122                    referenced_columns[source].update(columns)
123
124                column_aliases = node.alias_column_names
125                if column_aliases:
126                    source_column_alias_count[source] = len(column_aliases)
127
128    return expression
129
130
131def _remove_unused_selections(scope, parent_selections, schema, alias_count):
132    order = scope.expression.args.get("order")
133
134    if order:
135        # Assume columns without a qualified table are references to output columns
136        order_refs = {c.name for c in order.find_all(exp.Column) if not c.table}
137    else:
138        order_refs = set()
139
140    new_selections = []
141    removed = False
142    star = False
143    is_agg = False
144
145    select_all = SELECT_ALL in parent_selections
146
147    for selection in scope.expression.selects:
148        name = selection.alias_or_name
149
150        if select_all or name in parent_selections or name in order_refs or alias_count > 0:
151            new_selections.append(selection)
152            alias_count -= 1
153        else:
154            if selection.is_star:
155                star = True
156            removed = True
157
158        if not is_agg and selection.find(exp.AggFunc):
159            is_agg = True
160
161    if star:
162        resolver = Resolver(scope, schema)
163        names = {s.alias_or_name for s in new_selections}
164
165        for name in sorted(parent_selections):
166            if name not in names:
167                new_selections.append(
168                    alias(exp.column(name, table=resolver.get_table(name)), name, copy=False)
169                )
170
171    # If there are no remaining selections, just select a single constant
172    if not new_selections:
173        new_selections.append(default_selection(is_agg))
174
175    scope.expression.select(*new_selections, append=False, copy=False)
176
177    if removed:
178        scope.clear_cache()