sqlglot.optimizer.unnest_subqueries

  1from sqlglot import exp
  2from sqlglot.helper import name_sequence
  3from sqlglot.optimizer.scope import ScopeType, traverse_scope
  4
  5
  6def unnest_subqueries(expression):
  7    """
  8    Rewrite sqlglot AST to convert some predicates with subqueries into joins.
  9
 10    Convert scalar subqueries into cross joins.
 11    Convert correlated or vectorized subqueries into a group by so it is not a many to many left join.
 12
 13    Example:
 14        >>> import sqlglot
 15        >>> expression = sqlglot.parse_one("SELECT * FROM x AS x WHERE (SELECT y.a AS a FROM y AS y WHERE x.a = y.a) = 1 ")
 16        >>> unnest_subqueries(expression).sql()
 17        'SELECT * FROM x AS x LEFT JOIN (SELECT y.a AS a FROM y AS y WHERE TRUE GROUP BY y.a) AS _u_0 ON x.a = _u_0.a WHERE _u_0.a = 1'
 18
 19    Args:
 20        expression (sqlglot.Expression): expression to unnest
 21    Returns:
 22        sqlglot.Expression: unnested expression
 23    """
 24    next_alias_name = name_sequence("_u_")
 25
 26    for scope in traverse_scope(expression):
 27        select = scope.expression
 28        parent = select.parent_select
 29        if not parent:
 30            continue
 31        if scope.external_columns:
 32            decorrelate(select, parent, scope.external_columns, next_alias_name)
 33        elif scope.scope_type == ScopeType.SUBQUERY:
 34            unnest(select, parent, next_alias_name)
 35
 36    return expression
 37
 38
 39def unnest(select, parent_select, next_alias_name):
 40    if len(select.selects) > 1:
 41        return
 42
 43    predicate = select.find_ancestor(exp.Condition)
 44    alias = next_alias_name()
 45
 46    if not predicate or parent_select is not predicate.parent_select:
 47        return
 48
 49    # this subquery returns a scalar and can just be converted to a cross join
 50    if not isinstance(predicate, (exp.In, exp.Any)):
 51        having = predicate.find_ancestor(exp.Having)
 52        column = exp.column(select.selects[0].alias_or_name, alias)
 53        if having and having.parent_select is parent_select:
 54            column = exp.Max(this=column)
 55        _replace(select.parent, column)
 56
 57        parent_select.join(
 58            select,
 59            join_type="CROSS",
 60            join_alias=alias,
 61            copy=False,
 62        )
 63        return
 64
 65    if select.find(exp.Limit, exp.Offset):
 66        return
 67
 68    if isinstance(predicate, exp.Any):
 69        predicate = predicate.find_ancestor(exp.EQ)
 70
 71        if not predicate or parent_select is not predicate.parent_select:
 72            return
 73
 74    column = _other_operand(predicate)
 75    value = select.selects[0]
 76
 77    on = exp.condition(f'{column} = "{alias}"."{value.alias}"')
 78    _replace(predicate, f"NOT {on.right} IS NULL")
 79
 80    parent_select.join(
 81        select.group_by(value.this, copy=False),
 82        on=on,
 83        join_type="LEFT",
 84        join_alias=alias,
 85        copy=False,
 86    )
 87
 88
 89def decorrelate(select, parent_select, external_columns, next_alias_name):
 90    where = select.args.get("where")
 91
 92    if not where or where.find(exp.Or) or select.find(exp.Limit, exp.Offset):
 93        return
 94
 95    table_alias = next_alias_name()
 96    keys = []
 97
 98    # for all external columns in the where statement, find the relevant predicate
 99    # keys to convert it into a join
100    for column in external_columns:
101        if column.find_ancestor(exp.Where) is not where:
102            return
103
104        predicate = column.find_ancestor(exp.Predicate)
105
106        if not predicate or predicate.find_ancestor(exp.Where) is not where:
107            return
108
109        if isinstance(predicate, exp.Binary):
110            key = (
111                predicate.right
112                if any(node is column for node, *_ in predicate.left.walk())
113                else predicate.left
114            )
115        else:
116            return
117
118        keys.append((key, column, predicate))
119
120    if not any(isinstance(predicate, exp.EQ) for *_, predicate in keys):
121        return
122
123    is_subquery_projection = any(
124        node is select.parent for node in parent_select.selects if isinstance(node, exp.Subquery)
125    )
126
127    value = select.selects[0]
128    key_aliases = {}
129    group_by = []
130
131    for key, _, predicate in keys:
132        # if we filter on the value of the subquery, it needs to be unique
133        if key == value.this:
134            key_aliases[key] = value.alias
135            group_by.append(key)
136        else:
137            if key not in key_aliases:
138                key_aliases[key] = next_alias_name()
139            # all predicates that are equalities must also be in the unique
140            # so that we don't do a many to many join
141            if isinstance(predicate, exp.EQ) and key not in group_by:
142                group_by.append(key)
143
144    parent_predicate = select.find_ancestor(exp.Predicate)
145
146    # if the value of the subquery is not an agg or a key, we need to collect it into an array
147    # so that it can be grouped. For subquery projections, we use a MAX aggregation instead.
148    agg_func = exp.Max if is_subquery_projection else exp.ArrayAgg
149    if not value.find(exp.AggFunc) and value.this not in group_by:
150        select.select(
151            exp.alias_(agg_func(this=value.this), value.alias, quoted=False),
152            append=False,
153            copy=False,
154        )
155
156    # exists queries should not have any selects as it only checks if there are any rows
157    # all selects will be added by the optimizer and only used for join keys
158    if isinstance(parent_predicate, exp.Exists):
159        select.args["expressions"] = []
160
161    for key, alias in key_aliases.items():
162        if key in group_by:
163            # add all keys to the projections of the subquery
164            # so that we can use it as a join key
165            if isinstance(parent_predicate, exp.Exists) or key != value.this:
166                select.select(f"{key} AS {alias}", copy=False)
167        else:
168            select.select(exp.alias_(agg_func(this=key.copy()), alias, quoted=False), copy=False)
169
170    alias = exp.column(value.alias, table_alias)
171    other = _other_operand(parent_predicate)
172
173    if isinstance(parent_predicate, exp.Exists):
174        alias = exp.column(list(key_aliases.values())[0], table_alias)
175        parent_predicate = _replace(parent_predicate, f"NOT {alias} IS NULL")
176    elif isinstance(parent_predicate, exp.All):
177        parent_predicate = _replace(
178            parent_predicate.parent, f"ARRAY_ALL({alias}, _x -> _x = {other})"
179        )
180    elif isinstance(parent_predicate, exp.Any):
181        if value.this in group_by:
182            parent_predicate = _replace(parent_predicate.parent, f"{other} = {alias}")
183        else:
184            parent_predicate = _replace(parent_predicate, f"ARRAY_ANY({alias}, _x -> _x = {other})")
185    elif isinstance(parent_predicate, exp.In):
186        if value.this in group_by:
187            parent_predicate = _replace(parent_predicate, f"{other} = {alias}")
188        else:
189            parent_predicate = _replace(
190                parent_predicate,
191                f"ARRAY_ANY({alias}, _x -> _x = {parent_predicate.this})",
192            )
193    else:
194        if is_subquery_projection:
195            alias = exp.alias_(alias, select.parent.alias)
196
197        # COUNT always returns 0 on empty datasets, so we need take that into consideration here
198        # by transforming all counts into 0 and using that as the coalesced value
199        if value.find(exp.Count):
200
201            def remove_aggs(node):
202                if isinstance(node, exp.Count):
203                    return exp.Literal.number(0)
204                elif isinstance(node, exp.AggFunc):
205                    return exp.null()
206                return node
207
208            alias = exp.Coalesce(
209                this=alias,
210                expressions=[value.this.transform(remove_aggs)],
211            )
212
213        select.parent.replace(alias)
214
215    for key, column, predicate in keys:
216        predicate.replace(exp.true())
217        nested = exp.column(key_aliases[key], table_alias)
218
219        if is_subquery_projection:
220            key.replace(nested)
221            continue
222
223        if key in group_by:
224            key.replace(nested)
225        elif isinstance(predicate, exp.EQ):
226            parent_predicate = _replace(
227                parent_predicate,
228                f"({parent_predicate} AND ARRAY_CONTAINS({nested}, {column}))",
229            )
230        else:
231            key.replace(exp.to_identifier("_x"))
232            parent_predicate = _replace(
233                parent_predicate,
234                f'({parent_predicate} AND ARRAY_ANY({nested}, "_x" -> {predicate}))',
235            )
236
237    parent_select.join(
238        select.group_by(*group_by, copy=False),
239        on=[predicate for *_, predicate in keys if isinstance(predicate, exp.EQ)],
240        join_type="LEFT",
241        join_alias=table_alias,
242        copy=False,
243    )
244
245
246def _replace(expression, condition):
247    return expression.replace(exp.condition(condition))
248
249
250def _other_operand(expression):
251    if isinstance(expression, exp.In):
252        return expression.this
253
254    if isinstance(expression, (exp.Any, exp.All)):
255        return _other_operand(expression.parent)
256
257    if isinstance(expression, exp.Binary):
258        return (
259            expression.right
260            if isinstance(expression.left, (exp.Subquery, exp.Any, exp.Exists, exp.All))
261            else expression.left
262        )
263
264    return None