Edit on GitHub

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

Rewrite sqlglot AST to remove unused columns projections.

Example:
>>> import sqlglot
>>> sql = "SELECT y.a AS a FROM (SELECT x.a AS a, x.b AS b FROM x) AS y"
>>> expression = sqlglot.parse_one(sql)
>>> pushdown_projections(expression).sql()
'SELECT y.a AS a FROM (SELECT x.a AS a FROM x) AS y'
Arguments:
  • expression (sqlglot.Expression): expression to optimize
  • remove_unused_selections (bool): remove selects that are unused
Returns:

sqlglot.Expression: optimized expression