Edit on GitHub

sqlglot.optimizer.merge_subqueries

  1from collections import defaultdict
  2
  3from sqlglot import expressions as exp
  4from sqlglot.helper import find_new_name
  5from sqlglot.optimizer.scope import Scope, traverse_scope
  6
  7
  8def merge_subqueries(expression, leave_tables_isolated=False):
  9    """
 10    Rewrite sqlglot AST to merge derived tables into the outer query.
 11
 12    This also merges CTEs if they are selected from only once.
 13
 14    Example:
 15        >>> import sqlglot
 16        >>> expression = sqlglot.parse_one("SELECT a FROM (SELECT x.a FROM x) CROSS JOIN y")
 17        >>> merge_subqueries(expression).sql()
 18        'SELECT x.a FROM x CROSS JOIN y'
 19
 20    If `leave_tables_isolated` is True, this will not merge inner queries into outer
 21    queries if it would result in multiple table selects in a single query:
 22        >>> expression = sqlglot.parse_one("SELECT a FROM (SELECT x.a FROM x) CROSS JOIN y")
 23        >>> merge_subqueries(expression, leave_tables_isolated=True).sql()
 24        'SELECT a FROM (SELECT x.a FROM x) CROSS JOIN y'
 25
 26    Inspired by https://dev.mysql.com/doc/refman/8.0/en/derived-table-optimization.html
 27
 28    Args:
 29        expression (sqlglot.Expression): expression to optimize
 30        leave_tables_isolated (bool):
 31    Returns:
 32        sqlglot.Expression: optimized expression
 33    """
 34    expression = merge_ctes(expression, leave_tables_isolated)
 35    expression = merge_derived_tables(expression, leave_tables_isolated)
 36    return expression
 37
 38
 39# If a derived table has these Select args, it can't be merged
 40UNMERGABLE_ARGS = set(exp.Select.arg_types) - {
 41    "expressions",
 42    "from",
 43    "joins",
 44    "where",
 45    "order",
 46    "hint",
 47}
 48
 49
 50# Projections in the outer query that are instances of these types can be replaced
 51# without getting wrapped in parentheses, because the precedence won't be altered.
 52SAFE_TO_REPLACE_UNWRAPPED = (
 53    exp.Column,
 54    exp.EQ,
 55    exp.Func,
 56    exp.NEQ,
 57    exp.Paren,
 58)
 59
 60
 61def merge_ctes(expression, leave_tables_isolated=False):
 62    scopes = traverse_scope(expression)
 63
 64    # All places where we select from CTEs.
 65    # We key on the CTE scope so we can detect CTES that are selected from multiple times.
 66    cte_selections = defaultdict(list)
 67    for outer_scope in scopes:
 68        for table, inner_scope in outer_scope.selected_sources.values():
 69            if isinstance(inner_scope, Scope) and inner_scope.is_cte:
 70                cte_selections[id(inner_scope)].append(
 71                    (
 72                        outer_scope,
 73                        inner_scope,
 74                        table,
 75                    )
 76                )
 77
 78    singular_cte_selections = [v[0] for k, v in cte_selections.items() if len(v) == 1]
 79    for outer_scope, inner_scope, table in singular_cte_selections:
 80        from_or_join = table.find_ancestor(exp.From, exp.Join)
 81        if _mergeable(outer_scope, inner_scope, leave_tables_isolated, from_or_join):
 82            alias = table.alias_or_name
 83            _rename_inner_sources(outer_scope, inner_scope, alias)
 84            _merge_from(outer_scope, inner_scope, table, alias)
 85            _merge_expressions(outer_scope, inner_scope, alias)
 86            _merge_joins(outer_scope, inner_scope, from_or_join)
 87            _merge_where(outer_scope, inner_scope, from_or_join)
 88            _merge_order(outer_scope, inner_scope)
 89            _merge_hints(outer_scope, inner_scope)
 90            _pop_cte(inner_scope)
 91            outer_scope.clear_cache()
 92    return expression
 93
 94
 95def merge_derived_tables(expression, leave_tables_isolated=False):
 96    for outer_scope in traverse_scope(expression):
 97        for subquery in outer_scope.derived_tables:
 98            from_or_join = subquery.find_ancestor(exp.From, exp.Join)
 99            alias = subquery.alias_or_name
100            inner_scope = outer_scope.sources[alias]
101            if _mergeable(outer_scope, inner_scope, leave_tables_isolated, from_or_join):
102                _rename_inner_sources(outer_scope, inner_scope, alias)
103                _merge_from(outer_scope, inner_scope, subquery, alias)
104                _merge_expressions(outer_scope, inner_scope, alias)
105                _merge_joins(outer_scope, inner_scope, from_or_join)
106                _merge_where(outer_scope, inner_scope, from_or_join)
107                _merge_order(outer_scope, inner_scope)
108                _merge_hints(outer_scope, inner_scope)
109                outer_scope.clear_cache()
110
111    return expression
112
113
114def _mergeable(outer_scope, inner_scope, leave_tables_isolated, from_or_join):
115    """
116    Return True if `inner_select` can be merged into outer query.
117
118    Args:
119        outer_scope (Scope)
120        inner_scope (Scope)
121        leave_tables_isolated (bool)
122        from_or_join (exp.From|exp.Join)
123    Returns:
124        bool: True if can be merged
125    """
126    inner_select = inner_scope.expression.unnest()
127
128    def _is_a_window_expression_in_unmergable_operation():
129        window_expressions = inner_select.find_all(exp.Window)
130        window_alias_names = {window.parent.alias_or_name for window in window_expressions}
131        inner_select_name = from_or_join.alias_or_name
132        unmergable_window_columns = [
133            column
134            for column in outer_scope.columns
135            if column.find_ancestor(
136                exp.Where, exp.Group, exp.Order, exp.Join, exp.Having, exp.AggFunc
137            )
138        ]
139        window_expressions_in_unmergable = [
140            column
141            for column in unmergable_window_columns
142            if column.table == inner_select_name and column.name in window_alias_names
143        ]
144        return any(window_expressions_in_unmergable)
145
146    def _outer_select_joins_on_inner_select_join():
147        """
148        All columns from the inner select in the ON clause must be from the first FROM table.
149
150        That is, this can be merged:
151            SELECT * FROM x JOIN (SELECT y.a AS a FROM y JOIN z) AS q ON x.a = q.a
152                                         ^^^           ^
153        But this can't:
154            SELECT * FROM x JOIN (SELECT z.a AS a FROM y JOIN z) AS q ON x.a = q.a
155                                         ^^^                  ^
156        """
157        if not isinstance(from_or_join, exp.Join):
158            return False
159
160        alias = from_or_join.alias_or_name
161
162        on = from_or_join.args.get("on")
163        if not on:
164            return False
165        selections = [c.name for c in on.find_all(exp.Column) if c.table == alias]
166        inner_from = inner_scope.expression.args.get("from")
167        if not inner_from:
168            return False
169        inner_from_table = inner_from.alias_or_name
170        inner_projections = {s.alias_or_name: s for s in inner_scope.expression.selects}
171        return any(
172            col.table != inner_from_table
173            for selection in selections
174            for col in inner_projections[selection].find_all(exp.Column)
175        )
176
177    def _is_recursive():
178        # Recursive CTEs look like this:
179        #     WITH RECURSIVE cte AS (
180        #       SELECT * FROM x  <-- inner scope
181        #       UNION ALL
182        #       SELECT * FROM cte  <-- outer scope
183        #     )
184        cte = inner_scope.expression.parent
185        node = outer_scope.expression.parent
186
187        while node:
188            if node is cte:
189                return True
190            node = node.parent
191        return False
192
193    return (
194        isinstance(outer_scope.expression, exp.Select)
195        and not outer_scope.expression.is_star
196        and isinstance(inner_select, exp.Select)
197        and not any(inner_select.args.get(arg) for arg in UNMERGABLE_ARGS)
198        and inner_select.args.get("from")
199        and not outer_scope.pivots
200        and not any(e.find(exp.AggFunc, exp.Select, exp.Explode) for e in inner_select.expressions)
201        and not (leave_tables_isolated and len(outer_scope.selected_sources) > 1)
202        and not (
203            isinstance(from_or_join, exp.Join)
204            and inner_select.args.get("where")
205            and from_or_join.side in ("FULL", "LEFT", "RIGHT")
206        )
207        and not (
208            isinstance(from_or_join, exp.From)
209            and inner_select.args.get("where")
210            and any(
211                j.side in ("FULL", "RIGHT") for j in outer_scope.expression.args.get("joins", [])
212            )
213        )
214        and not _outer_select_joins_on_inner_select_join()
215        and not _is_a_window_expression_in_unmergable_operation()
216        and not _is_recursive()
217    )
218
219
220def _rename_inner_sources(outer_scope, inner_scope, alias):
221    """
222    Renames any sources in the inner query that conflict with names in the outer query.
223
224    Args:
225        outer_scope (sqlglot.optimizer.scope.Scope)
226        inner_scope (sqlglot.optimizer.scope.Scope)
227        alias (str)
228    """
229    taken = set(outer_scope.selected_sources)
230    conflicts = taken.intersection(set(inner_scope.selected_sources))
231    conflicts -= {alias}
232
233    for conflict in conflicts:
234        new_name = find_new_name(taken, conflict)
235
236        source, _ = inner_scope.selected_sources[conflict]
237        new_alias = exp.to_identifier(new_name)
238
239        if isinstance(source, exp.Subquery):
240            source.set("alias", exp.TableAlias(this=new_alias))
241        elif isinstance(source, exp.Table) and source.alias:
242            source.set("alias", new_alias)
243        elif isinstance(source, exp.Table):
244            source.replace(exp.alias_(source, new_alias))
245
246        for column in inner_scope.source_columns(conflict):
247            column.set("table", exp.to_identifier(new_name))
248
249        inner_scope.rename_source(conflict, new_name)
250
251
252def _merge_from(outer_scope, inner_scope, node_to_replace, alias):
253    """
254    Merge FROM clause of inner query into outer query.
255
256    Args:
257        outer_scope (sqlglot.optimizer.scope.Scope)
258        inner_scope (sqlglot.optimizer.scope.Scope)
259        node_to_replace (exp.Subquery|exp.Table)
260        alias (str)
261    """
262    new_subquery = inner_scope.expression.args["from"].this
263    new_subquery.set("joins", node_to_replace.args.get("joins"))
264    node_to_replace.replace(new_subquery)
265    for join_hint in outer_scope.join_hints:
266        tables = join_hint.find_all(exp.Table)
267        for table in tables:
268            if table.alias_or_name == node_to_replace.alias_or_name:
269                table.set("this", exp.to_identifier(new_subquery.alias_or_name))
270    outer_scope.remove_source(alias)
271    outer_scope.add_source(
272        new_subquery.alias_or_name, inner_scope.sources[new_subquery.alias_or_name]
273    )
274
275
276def _merge_joins(outer_scope, inner_scope, from_or_join):
277    """
278    Merge JOIN clauses of inner query into outer query.
279
280    Args:
281        outer_scope (sqlglot.optimizer.scope.Scope)
282        inner_scope (sqlglot.optimizer.scope.Scope)
283        from_or_join (exp.From|exp.Join)
284    """
285
286    new_joins = []
287
288    joins = inner_scope.expression.args.get("joins") or []
289    for join in joins:
290        new_joins.append(join)
291        outer_scope.add_source(join.alias_or_name, inner_scope.sources[join.alias_or_name])
292
293    if new_joins:
294        outer_joins = outer_scope.expression.args.get("joins", [])
295
296        # Maintain the join order
297        if isinstance(from_or_join, exp.From):
298            position = 0
299        else:
300            position = outer_joins.index(from_or_join) + 1
301        outer_joins[position:position] = new_joins
302
303        outer_scope.expression.set("joins", outer_joins)
304
305
306def _merge_expressions(outer_scope, inner_scope, alias):
307    """
308    Merge projections of inner query into outer query.
309
310    Args:
311        outer_scope (sqlglot.optimizer.scope.Scope)
312        inner_scope (sqlglot.optimizer.scope.Scope)
313        alias (str)
314    """
315    # Collect all columns that reference the alias of the inner query
316    outer_columns = defaultdict(list)
317    for column in outer_scope.columns:
318        if column.table == alias:
319            outer_columns[column.name].append(column)
320
321    # Replace columns with the projection expression in the inner query
322    for expression in inner_scope.expression.expressions:
323        projection_name = expression.alias_or_name
324        if not projection_name:
325            continue
326        columns_to_replace = outer_columns.get(projection_name, [])
327
328        expression = expression.unalias()
329        must_wrap_expression = not isinstance(expression, SAFE_TO_REPLACE_UNWRAPPED)
330
331        for column in columns_to_replace:
332            # Ensures we don't alter the intended operator precedence if there's additional
333            # context surrounding the outer expression (i.e. it's not a simple projection).
334            if isinstance(column.parent, (exp.Unary, exp.Binary)) and must_wrap_expression:
335                expression = exp.paren(expression, copy=False)
336
337            column.replace(expression.copy())
338
339
340def _merge_where(outer_scope, inner_scope, from_or_join):
341    """
342    Merge WHERE clause of inner query into outer query.
343
344    Args:
345        outer_scope (sqlglot.optimizer.scope.Scope)
346        inner_scope (sqlglot.optimizer.scope.Scope)
347        from_or_join (exp.From|exp.Join)
348    """
349    where = inner_scope.expression.args.get("where")
350    if not where or not where.this:
351        return
352
353    expression = outer_scope.expression
354
355    if isinstance(from_or_join, exp.Join):
356        # Merge predicates from an outer join to the ON clause
357        # if it only has columns that are already joined
358        from_ = expression.args.get("from")
359        sources = {from_.alias_or_name} if from_ else {}
360
361        for join in expression.args["joins"]:
362            source = join.alias_or_name
363            sources.add(source)
364            if source == from_or_join.alias_or_name:
365                break
366
367        if exp.column_table_names(where.this) <= sources:
368            from_or_join.on(where.this, copy=False)
369            from_or_join.set("on", from_or_join.args.get("on"))
370            return
371
372    expression.where(where.this, copy=False)
373
374
375def _merge_order(outer_scope, inner_scope):
376    """
377    Merge ORDER clause of inner query into outer query.
378
379    Args:
380        outer_scope (sqlglot.optimizer.scope.Scope)
381        inner_scope (sqlglot.optimizer.scope.Scope)
382    """
383    if (
384        any(
385            outer_scope.expression.args.get(arg) for arg in ["group", "distinct", "having", "order"]
386        )
387        or len(outer_scope.selected_sources) != 1
388        or any(expression.find(exp.AggFunc) for expression in outer_scope.expression.expressions)
389    ):
390        return
391
392    outer_scope.expression.set("order", inner_scope.expression.args.get("order"))
393
394
395def _merge_hints(outer_scope, inner_scope):
396    inner_scope_hint = inner_scope.expression.args.get("hint")
397    if not inner_scope_hint:
398        return
399    outer_scope_hint = outer_scope.expression.args.get("hint")
400    if outer_scope_hint:
401        for hint_expression in inner_scope_hint.expressions:
402            outer_scope_hint.append("expressions", hint_expression)
403    else:
404        outer_scope.expression.set("hint", inner_scope_hint)
405
406
407def _pop_cte(inner_scope):
408    """
409    Remove CTE from the AST.
410
411    Args:
412        inner_scope (sqlglot.optimizer.scope.Scope)
413    """
414    cte = inner_scope.expression.parent
415    with_ = cte.parent
416    if len(with_.expressions) == 1:
417        with_.pop()
418    else:
419        cte.pop()
def merge_subqueries(expression, leave_tables_isolated=False):
 9def merge_subqueries(expression, leave_tables_isolated=False):
10    """
11    Rewrite sqlglot AST to merge derived tables into the outer query.
12
13    This also merges CTEs if they are selected from only once.
14
15    Example:
16        >>> import sqlglot
17        >>> expression = sqlglot.parse_one("SELECT a FROM (SELECT x.a FROM x) CROSS JOIN y")
18        >>> merge_subqueries(expression).sql()
19        'SELECT x.a FROM x CROSS JOIN y'
20
21    If `leave_tables_isolated` is True, this will not merge inner queries into outer
22    queries if it would result in multiple table selects in a single query:
23        >>> expression = sqlglot.parse_one("SELECT a FROM (SELECT x.a FROM x) CROSS JOIN y")
24        >>> merge_subqueries(expression, leave_tables_isolated=True).sql()
25        'SELECT a FROM (SELECT x.a FROM x) CROSS JOIN y'
26
27    Inspired by https://dev.mysql.com/doc/refman/8.0/en/derived-table-optimization.html
28
29    Args:
30        expression (sqlglot.Expression): expression to optimize
31        leave_tables_isolated (bool):
32    Returns:
33        sqlglot.Expression: optimized expression
34    """
35    expression = merge_ctes(expression, leave_tables_isolated)
36    expression = merge_derived_tables(expression, leave_tables_isolated)
37    return expression

Rewrite sqlglot AST to merge derived tables into the outer query.

This also merges CTEs if they are selected from only once.

Example:
>>> import sqlglot
>>> expression = sqlglot.parse_one("SELECT a FROM (SELECT x.a FROM x) CROSS JOIN y")
>>> merge_subqueries(expression).sql()
'SELECT x.a FROM x CROSS JOIN y'

If leave_tables_isolated is True, this will not merge inner queries into outer queries if it would result in multiple table selects in a single query:

expression = sqlglot.parse_one("SELECT a FROM (SELECT x.a FROM x) CROSS JOIN y") merge_subqueries(expression, leave_tables_isolated=True).sql() 'SELECT a FROM (SELECT x.a FROM x) CROSS JOIN y'

Inspired by https://dev.mysql.com/doc/refman/8.0/en/derived-table-optimization.html

Arguments:
  • expression (sqlglot.Expression): expression to optimize
  • leave_tables_isolated (bool):
Returns:

sqlglot.Expression: optimized expression

UNMERGABLE_ARGS = {'with', 'match', 'cluster', 'windows', 'sample', 'distinct', 'prewhere', 'offset', 'laterals', 'distribute', 'into', 'locks', 'kind', 'group', 'format', 'having', 'settings', 'pivots', 'qualify', 'options', 'sort', 'connect', 'limit'}
SAFE_TO_REPLACE_UNWRAPPED = (<class 'sqlglot.expressions.Column'>, <class 'sqlglot.expressions.EQ'>, <class 'sqlglot.expressions.Func'>, <class 'sqlglot.expressions.NEQ'>, <class 'sqlglot.expressions.Paren'>)
def merge_ctes(expression, leave_tables_isolated=False):
62def merge_ctes(expression, leave_tables_isolated=False):
63    scopes = traverse_scope(expression)
64
65    # All places where we select from CTEs.
66    # We key on the CTE scope so we can detect CTES that are selected from multiple times.
67    cte_selections = defaultdict(list)
68    for outer_scope in scopes:
69        for table, inner_scope in outer_scope.selected_sources.values():
70            if isinstance(inner_scope, Scope) and inner_scope.is_cte:
71                cte_selections[id(inner_scope)].append(
72                    (
73                        outer_scope,
74                        inner_scope,
75                        table,
76                    )
77                )
78
79    singular_cte_selections = [v[0] for k, v in cte_selections.items() if len(v) == 1]
80    for outer_scope, inner_scope, table in singular_cte_selections:
81        from_or_join = table.find_ancestor(exp.From, exp.Join)
82        if _mergeable(outer_scope, inner_scope, leave_tables_isolated, from_or_join):
83            alias = table.alias_or_name
84            _rename_inner_sources(outer_scope, inner_scope, alias)
85            _merge_from(outer_scope, inner_scope, table, alias)
86            _merge_expressions(outer_scope, inner_scope, alias)
87            _merge_joins(outer_scope, inner_scope, from_or_join)
88            _merge_where(outer_scope, inner_scope, from_or_join)
89            _merge_order(outer_scope, inner_scope)
90            _merge_hints(outer_scope, inner_scope)
91            _pop_cte(inner_scope)
92            outer_scope.clear_cache()
93    return expression
def merge_derived_tables(expression, leave_tables_isolated=False):
 96def merge_derived_tables(expression, leave_tables_isolated=False):
 97    for outer_scope in traverse_scope(expression):
 98        for subquery in outer_scope.derived_tables:
 99            from_or_join = subquery.find_ancestor(exp.From, exp.Join)
100            alias = subquery.alias_or_name
101            inner_scope = outer_scope.sources[alias]
102            if _mergeable(outer_scope, inner_scope, leave_tables_isolated, from_or_join):
103                _rename_inner_sources(outer_scope, inner_scope, alias)
104                _merge_from(outer_scope, inner_scope, subquery, alias)
105                _merge_expressions(outer_scope, inner_scope, alias)
106                _merge_joins(outer_scope, inner_scope, from_or_join)
107                _merge_where(outer_scope, inner_scope, from_or_join)
108                _merge_order(outer_scope, inner_scope)
109                _merge_hints(outer_scope, inner_scope)
110                outer_scope.clear_cache()
111
112    return expression