Edit on GitHub

# sqlglot.optimizer.unnest_subqueries

View Source

      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    if (
     45        not predicate
     46        or parent_select is not predicate.parent_select
     47        or not parent_select.args.get("from")
     48    ):
     49        return
     50
     51    if isinstance(select, exp.Union):
     52        select = exp.select(*select.selects).from_(select.subquery(next_alias_name()))
     53
     54    alias = next_alias_name()
     55    clause = predicate.find_ancestor(exp.Having, exp.Where, exp.Join)
     56
     57    # This subquery returns a scalar and can just be converted to a cross join
     58    if not isinstance(predicate, (exp.In, exp.Any)):
     59        column = exp.column(select.selects[0].alias_or_name, alias)
     60
     61        clause_parent_select = clause.parent_select if clause else None
     62
     63        if (isinstance(clause, exp.Having) and clause_parent_select is parent_select) or (
     64            (not clause or clause_parent_select is not parent_select)
     65            and (
     66                parent_select.args.get("group")
     67                or any(projection.find(exp.AggFunc) for projection in parent_select.selects)
     68            )
     69        ):
     70            column = exp.Max(this=column)
     71        elif not isinstance(select.parent, exp.Subquery):
     72            return
     73
     74        _replace(select.parent, column)
     75        parent_select.join(select, join_type="CROSS", join_alias=alias, copy=False)
     76        return
     77
     78    if select.find(exp.Limit, exp.Offset):
     79        return
     80
     81    if isinstance(predicate, exp.Any):
     82        predicate = predicate.find_ancestor(exp.EQ)
     83
     84        if not predicate or parent_select is not predicate.parent_select:
     85            return
     86
     87    column = _other_operand(predicate)
     88    value = select.selects[0]
     89
     90    join_key = exp.column(value.alias, alias)
     91    join_key_not_null = join_key.is_(exp.null()).not_()
     92
     93    if isinstance(clause, exp.Join):
     94        _replace(predicate, exp.true())
     95        parent_select.where(join_key_not_null, copy=False)
     96    else:
     97        _replace(predicate, join_key_not_null)
     98
     99    group = select.args.get("group")
    100
    101    if group:
    102        if {value.this} != set(group.expressions):
    103            select = (
    104                exp.select(exp.column(value.alias, "_q"))
    105                .from_(select.subquery("_q", copy=False), copy=False)
    106                .group_by(exp.column(value.alias, "_q"), copy=False)
    107            )
    108    else:
    109        select = select.group_by(value.this, copy=False)
    110
    111    parent_select.join(
    112        select,
    113        on=column.eq(join_key),
    114        join_type="LEFT",
    115        join_alias=alias,
    116        copy=False,
    117    )
    118
    119
    120def decorrelate(select, parent_select, external_columns, next_alias_name):
    121    where = select.args.get("where")
    122
    123    if not where or where.find(exp.Or) or select.find(exp.Limit, exp.Offset):
    124        return
    125
    126    table_alias = next_alias_name()
    127    keys = []
    128
    129    # for all external columns in the where statement, find the relevant predicate
    130    # keys to convert it into a join
    131    for column in external_columns:
    132        if column.find_ancestor(exp.Where) is not where:
    133            return
    134
    135        predicate = column.find_ancestor(exp.Predicate)
    136
    137        if not predicate or predicate.find_ancestor(exp.Where) is not where:
    138            return
    139
    140        if isinstance(predicate, exp.Binary):
    141            key = (
    142                predicate.right
    143                if any(node is column for node in predicate.left.walk())
    144                else predicate.left
    145            )
    146        else:
    147            return
    148
    149        keys.append((key, column, predicate))
    150
    151    if not any(isinstance(predicate, exp.EQ) for *_, predicate in keys):
    152        return
    153
    154    is_subquery_projection = any(
    155        node is select.parent for node in parent_select.selects if isinstance(node, exp.Subquery)
    156    )
    157
    158    value = select.selects[0]
    159    key_aliases = {}
    160    group_by = []
    161
    162    for key, _, predicate in keys:
    163        # if we filter on the value of the subquery, it needs to be unique
    164        if key == value.this:
    165            key_aliases[key] = value.alias
    166            group_by.append(key)
    167        else:
    168            if key not in key_aliases:
    169                key_aliases[key] = next_alias_name()
    170            # all predicates that are equalities must also be in the unique
    171            # so that we don't do a many to many join
    172            if isinstance(predicate, exp.EQ) and key not in group_by:
    173                group_by.append(key)
    174
    175    parent_predicate = select.find_ancestor(exp.Predicate)
    176
    177    # if the value of the subquery is not an agg or a key, we need to collect it into an array
    178    # so that it can be grouped. For subquery projections, we use a MAX aggregation instead.
    179    agg_func = exp.Max if is_subquery_projection else exp.ArrayAgg
    180    if not value.find(exp.AggFunc) and value.this not in group_by:
    181        select.select(
    182            exp.alias_(agg_func(this=value.this), value.alias, quoted=False),
    183            append=False,
    184            copy=False,
    185        )
    186
    187    # exists queries should not have any selects as it only checks if there are any rows
    188    # all selects will be added by the optimizer and only used for join keys
    189    if isinstance(parent_predicate, exp.Exists):
    190        select.args["expressions"] = []
    191
    192    for key, alias in key_aliases.items():
    193        if key in group_by:
    194            # add all keys to the projections of the subquery
    195            # so that we can use it as a join key
    196            if isinstance(parent_predicate, exp.Exists) or key != value.this:
    197                select.select(f"{key} AS {alias}", copy=False)
    198        else:
    199            select.select(exp.alias_(agg_func(this=key.copy()), alias, quoted=False), copy=False)
    200
    201    alias = exp.column(value.alias, table_alias)
    202    other = _other_operand(parent_predicate)
    203
    204    if isinstance(parent_predicate, exp.Exists):
    205        alias = exp.column(list(key_aliases.values())[0], table_alias)
    206        parent_predicate = _replace(parent_predicate, f"NOT {alias} IS NULL")
    207    elif isinstance(parent_predicate, exp.All):
    208        parent_predicate = _replace(
    209            parent_predicate.parent, f"ARRAY_ALL({alias}, _x -> _x = {other})"
    210        )
    211    elif isinstance(parent_predicate, exp.Any):
    212        if value.this in group_by:
    213            parent_predicate = _replace(parent_predicate.parent, f"{other} = {alias}")
    214        else:
    215            parent_predicate = _replace(parent_predicate, f"ARRAY_ANY({alias}, _x -> _x = {other})")
    216    elif isinstance(parent_predicate, exp.In):
    217        if value.this in group_by:
    218            parent_predicate = _replace(parent_predicate, f"{other} = {alias}")
    219        else:
    220            parent_predicate = _replace(
    221                parent_predicate,
    222                f"ARRAY_ANY({alias}, _x -> _x = {parent_predicate.this})",
    223            )
    224    else:
    225        if is_subquery_projection:
    226            alias = exp.alias_(alias, select.parent.alias)
    227
    228        # COUNT always returns 0 on empty datasets, so we need take that into consideration here
    229        # by transforming all counts into 0 and using that as the coalesced value
    230        if value.find(exp.Count):
    231
    232            def remove_aggs(node):
    233                if isinstance(node, exp.Count):
    234                    return exp.Literal.number(0)
    235                elif isinstance(node, exp.AggFunc):
    236                    return exp.null()
    237                return node
    238
    239            alias = exp.Coalesce(
    240                this=alias,
    241                expressions=[value.this.transform(remove_aggs)],
    242            )
    243
    244        select.parent.replace(alias)
    245
    246    for key, column, predicate in keys:
    247        predicate.replace(exp.true())
    248        nested = exp.column(key_aliases[key], table_alias)
    249
    250        if is_subquery_projection:
    251            key.replace(nested)
    252            continue
    253
    254        if key in group_by:
    255            key.replace(nested)
    256        elif isinstance(predicate, exp.EQ):
    257            parent_predicate = _replace(
    258                parent_predicate,
    259                f"({parent_predicate} AND ARRAY_CONTAINS({nested}, {column}))",
    260            )
    261        else:
    262            key.replace(exp.to_identifier("_x"))
    263            parent_predicate = _replace(
    264                parent_predicate,
    265                f"({parent_predicate} AND ARRAY_ANY({nested}, _x -> {predicate}))",
    266            )
    267
    268    parent_select.join(
    269        select.group_by(*group_by, copy=False),
    270        on=[predicate for *_, predicate in keys if isinstance(predicate, exp.EQ)],
    271        join_type="LEFT",
    272        join_alias=table_alias,
    273        copy=False,
    274    )
    275
    276
    277def _replace(expression, condition):
    278    return expression.replace(exp.condition(condition))
    279
    280
    281def _other_operand(expression):
    282    if isinstance(expression, exp.In):
    283        return expression.this
    284
    285    if isinstance(expression, (exp.Any, exp.All)):
    286        return _other_operand(expression.parent)
    287
    288    if isinstance(expression, exp.Binary):
    289        return (
    290            expression.right
    291            if isinstance(expression.left, (exp.Subquery, exp.Any, exp.Exists, exp.All))
    292            else expression.left
    293        )
    294
    295    return None

def unnest_subqueries(expression): View Source

     7def unnest_subqueries(expression):
     8    """
     9    Rewrite sqlglot AST to convert some predicates with subqueries into joins.
    10
    11    Convert scalar subqueries into cross joins.
    12    Convert correlated or vectorized subqueries into a group by so it is not a many to many left join.
    13
    14    Example:
    15        >>> import sqlglot
    16        >>> 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 ")
    17        >>> unnest_subqueries(expression).sql()
    18        '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'
    19
    20    Args:
    21        expression (sqlglot.Expression): expression to unnest
    22    Returns:
    23        sqlglot.Expression: unnested expression
    24    """
    25    next_alias_name = name_sequence("_u_")
    26
    27    for scope in traverse_scope(expression):
    28        select = scope.expression
    29        parent = select.parent_select
    30        if not parent:
    31            continue
    32        if scope.external_columns:
    33            decorrelate(select, parent, scope.external_columns, next_alias_name)
    34        elif scope.scope_type == ScopeType.SUBQUERY:
    35            unnest(select, parent, next_alias_name)
    36
    37    return expression

Rewrite sqlglot AST to convert some predicates with subqueries into joins.

Convert scalar subqueries into cross joins. Convert correlated or vectorized subqueries into a group by so it is not a many to many left join.

###### Example:

>     >>> import sqlglot
>     >>> 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 ")
>     >>> unnest_subqueries(expression).sql()
>     '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'

###### Arguments:

- expression (sqlglot.Expression): expression to unnest

###### Returns:

> sqlglot.Expression: unnested expression

def unnest(select, parent_select, next_alias_name): View Source

     40def unnest(select, parent_select, next_alias_name):
     41    if len(select.selects) > 1:
     42        return
     43
     44    predicate = select.find_ancestor(exp.Condition)
     45    if (
     46        not predicate
     47        or parent_select is not predicate.parent_select
     48        or not parent_select.args.get("from")
     49    ):
     50        return
     51
     52    if isinstance(select, exp.Union):
     53        select = exp.select(*select.selects).from_(select.subquery(next_alias_name()))
     54
     55    alias = next_alias_name()
     56    clause = predicate.find_ancestor(exp.Having, exp.Where, exp.Join)
     57
     58    # This subquery returns a scalar and can just be converted to a cross join
     59    if not isinstance(predicate, (exp.In, exp.Any)):
     60        column = exp.column(select.selects[0].alias_or_name, alias)
     61
     62        clause_parent_select = clause.parent_select if clause else None
     63
     64        if (isinstance(clause, exp.Having) and clause_parent_select is parent_select) or (
     65            (not clause or clause_parent_select is not parent_select)
     66            and (
     67                parent_select.args.get("group")
     68                or any(projection.find(exp.AggFunc) for projection in parent_select.selects)
     69            )
     70        ):
     71            column = exp.Max(this=column)
     72        elif not isinstance(select.parent, exp.Subquery):
     73            return
     74
     75        _replace(select.parent, column)
     76        parent_select.join(select, join_type="CROSS", join_alias=alias, copy=False)
     77        return
     78
     79    if select.find(exp.Limit, exp.Offset):
     80        return
     81
     82    if isinstance(predicate, exp.Any):
     83        predicate = predicate.find_ancestor(exp.EQ)
     84
     85        if not predicate or parent_select is not predicate.parent_select:
     86            return
     87
     88    column = _other_operand(predicate)
     89    value = select.selects[0]
     90
     91    join_key = exp.column(value.alias, alias)
     92    join_key_not_null = join_key.is_(exp.null()).not_()
     93
     94    if isinstance(clause, exp.Join):
     95        _replace(predicate, exp.true())
     96        parent_select.where(join_key_not_null, copy=False)
     97    else:
     98        _replace(predicate, join_key_not_null)
     99
    100    group = select.args.get("group")
    101
    102    if group:
    103        if {value.this} != set(group.expressions):
    104            select = (
    105                exp.select(exp.column(value.alias, "_q"))
    106                .from_(select.subquery("_q", copy=False), copy=False)
    107                .group_by(exp.column(value.alias, "_q"), copy=False)
    108            )
    109    else:
    110        select = select.group_by(value.this, copy=False)
    111
    112    parent_select.join(
    113        select,
    114        on=column.eq(join_key),
    115        join_type="LEFT",
    116        join_alias=alias,
    117        copy=False,
    118    )

def decorrelate(select, parent_select, external_columns, next_alias_name): View Source

    121def decorrelate(select, parent_select, external_columns, next_alias_name):
    122    where = select.args.get("where")
    123
    124    if not where or where.find(exp.Or) or select.find(exp.Limit, exp.Offset):
    125        return
    126
    127    table_alias = next_alias_name()
    128    keys = []
    129
    130    # for all external columns in the where statement, find the relevant predicate
    131    # keys to convert it into a join
    132    for column in external_columns:
    133        if column.find_ancestor(exp.Where) is not where:
    134            return
    135
    136        predicate = column.find_ancestor(exp.Predicate)
    137
    138        if not predicate or predicate.find_ancestor(exp.Where) is not where:
    139            return
    140
    141        if isinstance(predicate, exp.Binary):
    142            key = (
    143                predicate.right
    144                if any(node is column for node in predicate.left.walk())
    145                else predicate.left
    146            )
    147        else:
    148            return
    149
    150        keys.append((key, column, predicate))
    151
    152    if not any(isinstance(predicate, exp.EQ) for *_, predicate in keys):
    153        return
    154
    155    is_subquery_projection = any(
    156        node is select.parent for node in parent_select.selects if isinstance(node, exp.Subquery)
    157    )
    158
    159    value = select.selects[0]
    160    key_aliases = {}
    161    group_by = []
    162
    163    for key, _, predicate in keys:
    164        # if we filter on the value of the subquery, it needs to be unique
    165        if key == value.this:
    166            key_aliases[key] = value.alias
    167            group_by.append(key)
    168        else:
    169            if key not in key_aliases:
    170                key_aliases[key] = next_alias_name()
    171            # all predicates that are equalities must also be in the unique
    172            # so that we don't do a many to many join
    173            if isinstance(predicate, exp.EQ) and key not in group_by:
    174                group_by.append(key)
    175
    176    parent_predicate = select.find_ancestor(exp.Predicate)
    177
    178    # if the value of the subquery is not an agg or a key, we need to collect it into an array
    179    # so that it can be grouped. For subquery projections, we use a MAX aggregation instead.
    180    agg_func = exp.Max if is_subquery_projection else exp.ArrayAgg
    181    if not value.find(exp.AggFunc) and value.this not in group_by:
    182        select.select(
    183            exp.alias_(agg_func(this=value.this), value.alias, quoted=False),
    184            append=False,
    185            copy=False,
    186        )
    187
    188    # exists queries should not have any selects as it only checks if there are any rows
    189    # all selects will be added by the optimizer and only used for join keys
    190    if isinstance(parent_predicate, exp.Exists):
    191        select.args["expressions"] = []
    192
    193    for key, alias in key_aliases.items():
    194        if key in group_by:
    195            # add all keys to the projections of the subquery
    196            # so that we can use it as a join key
    197            if isinstance(parent_predicate, exp.Exists) or key != value.this:
    198                select.select(f"{key} AS {alias}", copy=False)
    199        else:
    200            select.select(exp.alias_(agg_func(this=key.copy()), alias, quoted=False), copy=False)
    201
    202    alias = exp.column(value.alias, table_alias)
    203    other = _other_operand(parent_predicate)
    204
    205    if isinstance(parent_predicate, exp.Exists):
    206        alias = exp.column(list(key_aliases.values())[0], table_alias)
    207        parent_predicate = _replace(parent_predicate, f"NOT {alias} IS NULL")
    208    elif isinstance(parent_predicate, exp.All):
    209        parent_predicate = _replace(
    210            parent_predicate.parent, f"ARRAY_ALL({alias}, _x -> _x = {other})"
    211        )
    212    elif isinstance(parent_predicate, exp.Any):
    213        if value.this in group_by:
    214            parent_predicate = _replace(parent_predicate.parent, f"{other} = {alias}")
    215        else:
    216            parent_predicate = _replace(parent_predicate, f"ARRAY_ANY({alias}, _x -> _x = {other})")
    217    elif isinstance(parent_predicate, exp.In):
    218        if value.this in group_by:
    219            parent_predicate = _replace(parent_predicate, f"{other} = {alias}")
    220        else:
    221            parent_predicate = _replace(
    222                parent_predicate,
    223                f"ARRAY_ANY({alias}, _x -> _x = {parent_predicate.this})",
    224            )
    225    else:
    226        if is_subquery_projection:
    227            alias = exp.alias_(alias, select.parent.alias)
    228
    229        # COUNT always returns 0 on empty datasets, so we need take that into consideration here
    230        # by transforming all counts into 0 and using that as the coalesced value
    231        if value.find(exp.Count):
    232
    233            def remove_aggs(node):
    234                if isinstance(node, exp.Count):
    235                    return exp.Literal.number(0)
    236                elif isinstance(node, exp.AggFunc):
    237                    return exp.null()
    238                return node
    239
    240            alias = exp.Coalesce(
    241                this=alias,
    242                expressions=[value.this.transform(remove_aggs)],
    243            )
    244
    245        select.parent.replace(alias)
    246
    247    for key, column, predicate in keys:
    248        predicate.replace(exp.true())
    249        nested = exp.column(key_aliases[key], table_alias)
    250
    251        if is_subquery_projection:
    252            key.replace(nested)
    253            continue
    254
    255        if key in group_by:
    256            key.replace(nested)
    257        elif isinstance(predicate, exp.EQ):
    258            parent_predicate = _replace(
    259                parent_predicate,
    260                f"({parent_predicate} AND ARRAY_CONTAINS({nested}, {column}))",
    261            )
    262        else:
    263            key.replace(exp.to_identifier("_x"))
    264            parent_predicate = _replace(
    265                parent_predicate,
    266                f"({parent_predicate} AND ARRAY_ANY({nested}, _x -> {predicate}))",
    267            )
    268
    269    parent_select.join(
    270        select.group_by(*group_by, copy=False),
    271        on=[predicate for *_, predicate in keys if isinstance(predicate, exp.EQ)],
    272        join_type="LEFT",
    273        join_alias=table_alias,
    274        copy=False,
    275    )