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# Selection to use if selection list is empty 12DEFAULT_SELECTION = lambda is_agg: alias( 13 exp.Max(this=exp.Literal.number(1)) if is_agg else "1", "_" 14) 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 if scope.expression.args.get("distinct") or (scope.parent and scope.parent.pivots): 47 # We can't remove columns SELECT DISTINCT nor UNION DISTINCT. The same holds if 48 # we select from a pivoted source in the parent scope. 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 = 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 89 90 91def _remove_unused_selections(scope, parent_selections, schema, alias_count): 92 order = scope.expression.args.get("order") 93 94 if order: 95 # Assume columns without a qualified table are references to output columns 96 order_refs = {c.name for c in order.find_all(exp.Column) if not c.table} 97 else: 98 order_refs = set() 99 100 new_selections = [] 101 removed = False 102 star = False 103 is_agg = False 104 105 select_all = SELECT_ALL in parent_selections 106 107 for selection in scope.expression.selects: 108 name = selection.alias_or_name 109 110 if select_all or name in parent_selections or name in order_refs or alias_count > 0: 111 new_selections.append(selection) 112 alias_count -= 1 113 else: 114 if selection.is_star: 115 star = True 116 removed = True 117 118 if not is_agg and selection.find(exp.AggFunc): 119 is_agg = True 120 121 if star: 122 resolver = Resolver(scope, schema) 123 names = {s.alias_or_name for s in new_selections} 124 125 for name in sorted(parent_selections): 126 if name not in names: 127 new_selections.append( 128 alias(exp.column(name, table=resolver.get_table(name)), name, copy=False) 129 ) 130 131 # If there are no remaining selections, just select a single constant 132 if not new_selections: 133 new_selections.append(DEFAULT_SELECTION(is_agg)) 134 135 scope.expression.select(*new_selections, append=False, copy=False) 136 137 if removed: 138 scope.clear_cache()
SELECT_ALL =
<object object>
def
DEFAULT_SELECTION(is_agg):
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 if scope.expression.args.get("distinct") or (scope.parent and scope.parent.pivots): 48 # We can't remove columns SELECT DISTINCT nor UNION DISTINCT. The same holds if 49 # we select from a pivoted source in the parent scope. 50 parent_selections = {SELECT_ALL} 51 52 if isinstance(scope.expression, exp.Union): 53 left, right = scope.union_scopes 54 referenced_columns[left] = parent_selections 55 56 if any(select.is_star for select in right.expression.selects): 57 referenced_columns[right] = parent_selections 58 elif not any(select.is_star for select in left.expression.selects): 59 referenced_columns[right] = [ 60 right.expression.selects[i].alias_or_name 61 for i, select in enumerate(left.expression.selects) 62 if SELECT_ALL in parent_selections or select.alias_or_name in parent_selections 63 ] 64 65 if isinstance(scope.expression, exp.Select): 66 if remove_unused_selections: 67 _remove_unused_selections(scope, parent_selections, schema, alias_count) 68 69 if scope.expression.is_star: 70 continue 71 72 # Group columns by source name 73 selects = defaultdict(set) 74 for col in scope.columns: 75 table_name = col.table 76 col_name = col.name 77 selects[table_name].add(col_name) 78 79 # Push the selected columns down to the next scope 80 for name, (node, source) in scope.selected_sources.items(): 81 if isinstance(source, Scope): 82 columns = selects.get(name) or set() 83 referenced_columns[source].update(columns) 84 85 column_aliases = node.alias_column_names 86 if column_aliases: 87 source_column_alias_count[source] = len(column_aliases) 88 89 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