sqlglot.optimizer.pushdown_projections

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