Edit on GitHub

# sqlglot.transforms

View Source

      1from __future__ import annotations
      2
      3import typing as t
      4
      5from sqlglot import expressions as exp
      6from sqlglot.helper import find_new_name
      7
      8if t.TYPE_CHECKING:
      9    from sqlglot.generator import Generator
     10
     11
     12def unalias_group(expression: exp.Expression) -> exp.Expression:
     13    """
     14    Replace references to select aliases in GROUP BY clauses.
     15
     16    Example:
     17        >>> import sqlglot
     18        >>> sqlglot.parse_one("SELECT a AS b FROM x GROUP BY b").transform(unalias_group).sql()
     19        'SELECT a AS b FROM x GROUP BY 1'
     20
     21    Args:
     22        expression: the expression that will be transformed.
     23
     24    Returns:
     25        The transformed expression.
     26    """
     27    if isinstance(expression, exp.Group) and isinstance(expression.parent, exp.Select):
     28        aliased_selects = {
     29            e.alias: i
     30            for i, e in enumerate(expression.parent.expressions, start=1)
     31            if isinstance(e, exp.Alias)
     32        }
     33
     34        for group_by in expression.expressions:
     35            if (
     36                isinstance(group_by, exp.Column)
     37                and not group_by.table
     38                and group_by.name in aliased_selects
     39            ):
     40                group_by.replace(exp.Literal.number(aliased_selects.get(group_by.name)))
     41
     42    return expression
     43
     44
     45def eliminate_distinct_on(expression: exp.Expression) -> exp.Expression:
     46    """
     47    Convert SELECT DISTINCT ON statements to a subquery with a window function.
     48
     49    This is useful for dialects that don't support SELECT DISTINCT ON but support window functions.
     50
     51    Args:
     52        expression: the expression that will be transformed.
     53
     54    Returns:
     55        The transformed expression.
     56    """
     57    if (
     58        isinstance(expression, exp.Select)
     59        and expression.args.get("distinct")
     60        and expression.args["distinct"].args.get("on")
     61        and isinstance(expression.args["distinct"].args["on"], exp.Tuple)
     62    ):
     63        distinct_cols = expression.args["distinct"].pop().args["on"].expressions
     64        outer_selects = expression.selects
     65        row_number = find_new_name(expression.named_selects, "_row_number")
     66        window = exp.Window(
     67            this=exp.RowNumber(),
     68            partition_by=distinct_cols,
     69        )
     70        order = expression.args.get("order")
     71        if order:
     72            window.set("order", order.pop().copy())
     73        window = exp.alias_(window, row_number)
     74        expression.select(window, copy=False)
     75        return exp.select(*outer_selects).from_(expression.subquery()).where(f'"{row_number}" = 1')
     76    return expression
     77
     78
     79def eliminate_qualify(expression: exp.Expression) -> exp.Expression:
     80    """
     81    Convert SELECT statements that contain the QUALIFY clause into subqueries, filtered equivalently.
     82
     83    The idea behind this transformation can be seen in Snowflake's documentation for QUALIFY:
     84    https://docs.snowflake.com/en/sql-reference/constructs/qualify
     85
     86    Some dialects don't support window functions in the WHERE clause, so we need to include them as
     87    projections in the subquery, in order to refer to them in the outer filter using aliases. Also,
     88    if a column is referenced in the QUALIFY clause but is not selected, we need to include it too,
     89    otherwise we won't be able to refer to it in the outer query's WHERE clause.
     90    """
     91    if isinstance(expression, exp.Select) and expression.args.get("qualify"):
     92        taken = set(expression.named_selects)
     93        for select in expression.selects:
     94            if not select.alias_or_name:
     95                alias = find_new_name(taken, "_c")
     96                select.replace(exp.alias_(select.copy(), alias))
     97                taken.add(alias)
     98
     99        outer_selects = exp.select(*[select.alias_or_name for select in expression.selects])
    100        qualify_filters = expression.args["qualify"].pop().this
    101
    102        for expr in qualify_filters.find_all((exp.Window, exp.Column)):
    103            if isinstance(expr, exp.Window):
    104                alias = find_new_name(expression.named_selects, "_w")
    105                expression.select(exp.alias_(expr.copy(), alias), copy=False)
    106                column = exp.column(alias)
    107                if isinstance(expr.parent, exp.Qualify):
    108                    qualify_filters = column
    109                else:
    110                    expr.replace(column)
    111            elif expr.name not in expression.named_selects:
    112                expression.select(expr.copy(), copy=False)
    113
    114        return outer_selects.from_(expression.subquery(alias="_t")).where(qualify_filters)
    115
    116    return expression
    117
    118
    119def remove_precision_parameterized_types(expression: exp.Expression) -> exp.Expression:
    120    """
    121    Some dialects only allow the precision for parameterized types to be defined in the DDL and not in
    122    other expressions. This transforms removes the precision from parameterized types in expressions.
    123    """
    124    for node in expression.find_all(exp.DataType):
    125        node.set("expressions", [e for e in node.expressions if isinstance(e, exp.DataType)])
    126    return expression
    127
    128
    129def unnest_to_explode(expression: exp.Expression) -> exp.Expression:
    130    """Convert cross join unnest into lateral view explode (used in presto -> hive)."""
    131    if isinstance(expression, exp.Select):
    132        for join in expression.args.get("joins") or []:
    133            unnest = join.this
    134
    135            if isinstance(unnest, exp.Unnest):
    136                alias = unnest.args.get("alias")
    137                udtf = exp.Posexplode if unnest.args.get("ordinality") else exp.Explode
    138
    139                expression.args["joins"].remove(join)
    140
    141                for e, column in zip(unnest.expressions, alias.columns if alias else []):
    142                    expression.append(
    143                        "laterals",
    144                        exp.Lateral(
    145                            this=udtf(this=e),
    146                            view=True,
    147                            alias=exp.TableAlias(this=alias.this, columns=[column]),  # type: ignore
    148                        ),
    149                    )
    150    return expression
    151
    152
    153def explode_to_unnest(expression: exp.Expression) -> exp.Expression:
    154    """Convert explode/posexplode into unnest (used in hive -> presto)."""
    155    if isinstance(expression, exp.Select):
    156        from sqlglot.optimizer.scope import build_scope
    157
    158        taken_select_names = set(expression.named_selects)
    159        taken_source_names = set(build_scope(expression).selected_sources)
    160
    161        for select in expression.selects:
    162            to_replace = select
    163
    164            pos_alias = ""
    165            explode_alias = ""
    166
    167            if isinstance(select, exp.Alias):
    168                explode_alias = select.alias
    169                select = select.this
    170            elif isinstance(select, exp.Aliases):
    171                pos_alias = select.aliases[0].name
    172                explode_alias = select.aliases[1].name
    173                select = select.this
    174
    175            if isinstance(select, (exp.Explode, exp.Posexplode)):
    176                is_posexplode = isinstance(select, exp.Posexplode)
    177
    178                explode_arg = select.this
    179                unnest = exp.Unnest(expressions=[explode_arg.copy()], ordinality=is_posexplode)
    180
    181                # This ensures that we won't use [POS]EXPLODE's argument as a new selection
    182                if isinstance(explode_arg, exp.Column):
    183                    taken_select_names.add(explode_arg.output_name)
    184
    185                unnest_source_alias = find_new_name(taken_source_names, "_u")
    186                taken_source_names.add(unnest_source_alias)
    187
    188                if not explode_alias:
    189                    explode_alias = find_new_name(taken_select_names, "col")
    190                    taken_select_names.add(explode_alias)
    191
    192                    if is_posexplode:
    193                        pos_alias = find_new_name(taken_select_names, "pos")
    194                        taken_select_names.add(pos_alias)
    195
    196                if is_posexplode:
    197                    column_names = [explode_alias, pos_alias]
    198                    to_replace.pop()
    199                    expression.select(pos_alias, explode_alias, copy=False)
    200                else:
    201                    column_names = [explode_alias]
    202                    to_replace.replace(exp.column(explode_alias))
    203
    204                unnest = exp.alias_(unnest, unnest_source_alias, table=column_names)
    205
    206                if not expression.args.get("from"):
    207                    expression.from_(unnest, copy=False)
    208                else:
    209                    expression.join(unnest, join_type="CROSS", copy=False)
    210
    211    return expression
    212
    213
    214def remove_target_from_merge(expression: exp.Expression) -> exp.Expression:
    215    """Remove table refs from columns in when statements."""
    216    if isinstance(expression, exp.Merge):
    217        alias = expression.this.args.get("alias")
    218        targets = {expression.this.this}
    219        if alias:
    220            targets.add(alias.this)
    221
    222        for when in expression.expressions:
    223            when.transform(
    224                lambda node: exp.column(node.name)
    225                if isinstance(node, exp.Column) and node.args.get("table") in targets
    226                else node,
    227                copy=False,
    228            )
    229    return expression
    230
    231
    232def preprocess(
    233    transforms: t.List[t.Callable[[exp.Expression], exp.Expression]],
    234) -> t.Callable[[Generator, exp.Expression], str]:
    235    """
    236    Creates a new transform by chaining a sequence of transformations and converts the resulting
    237    expression to SQL, using an appropriate `Generator.TRANSFORMS` function.
    238
    239    Args:
    240        transforms: sequence of transform functions. These will be called in order.
    241
    242    Returns:
    243        Function that can be used as a generator transform.
    244    """
    245
    246    def _to_sql(self, expression: exp.Expression) -> str:
    247        expression = transforms[0](expression.copy())
    248        for t in transforms[1:]:
    249            expression = t(expression)
    250        return getattr(self, expression.key + "_sql")(expression)
    251
    252    return _to_sql

def unalias_group( expression: sqlglot.expressions.Expression) -> sqlglot.expressions.Expression: View Source

    13def unalias_group(expression: exp.Expression) -> exp.Expression:
    14    """
    15    Replace references to select aliases in GROUP BY clauses.
    16
    17    Example:
    18        >>> import sqlglot
    19        >>> sqlglot.parse_one("SELECT a AS b FROM x GROUP BY b").transform(unalias_group).sql()
    20        'SELECT a AS b FROM x GROUP BY 1'
    21
    22    Args:
    23        expression: the expression that will be transformed.
    24
    25    Returns:
    26        The transformed expression.
    27    """
    28    if isinstance(expression, exp.Group) and isinstance(expression.parent, exp.Select):
    29        aliased_selects = {
    30            e.alias: i
    31            for i, e in enumerate(expression.parent.expressions, start=1)
    32            if isinstance(e, exp.Alias)
    33        }
    34
    35        for group_by in expression.expressions:
    36            if (
    37                isinstance(group_by, exp.Column)
    38                and not group_by.table
    39                and group_by.name in aliased_selects
    40            ):
    41                group_by.replace(exp.Literal.number(aliased_selects.get(group_by.name)))
    42
    43    return expression

Replace references to select aliases in GROUP BY clauses.

###### Example:

>     >>> import sqlglot
>     >>> sqlglot.parse_one("SELECT a AS b FROM x GROUP BY b").transform(unalias_group).sql()
>     'SELECT a AS b FROM x GROUP BY 1'

###### Arguments:

- expression: the expression that will be transformed.

###### Returns:

> The transformed expression.

def eliminate_distinct_on( expression: sqlglot.expressions.Expression) -> sqlglot.expressions.Expression: View Source

    46def eliminate_distinct_on(expression: exp.Expression) -> exp.Expression:
    47    """
    48    Convert SELECT DISTINCT ON statements to a subquery with a window function.
    49
    50    This is useful for dialects that don't support SELECT DISTINCT ON but support window functions.
    51
    52    Args:
    53        expression: the expression that will be transformed.
    54
    55    Returns:
    56        The transformed expression.
    57    """
    58    if (
    59        isinstance(expression, exp.Select)
    60        and expression.args.get("distinct")
    61        and expression.args["distinct"].args.get("on")
    62        and isinstance(expression.args["distinct"].args["on"], exp.Tuple)
    63    ):
    64        distinct_cols = expression.args["distinct"].pop().args["on"].expressions
    65        outer_selects = expression.selects
    66        row_number = find_new_name(expression.named_selects, "_row_number")
    67        window = exp.Window(
    68            this=exp.RowNumber(),
    69            partition_by=distinct_cols,
    70        )
    71        order = expression.args.get("order")
    72        if order:
    73            window.set("order", order.pop().copy())
    74        window = exp.alias_(window, row_number)
    75        expression.select(window, copy=False)
    76        return exp.select(*outer_selects).from_(expression.subquery()).where(f'"{row_number}" = 1')
    77    return expression

Convert SELECT DISTINCT ON statements to a subquery with a window function.

This is useful for dialects that don't support SELECT DISTINCT ON but support window functions.

###### Arguments:

- expression: the expression that will be transformed.

###### Returns:

> The transformed expression.

def eliminate_qualify( expression: sqlglot.expressions.Expression) -> sqlglot.expressions.Expression: View Source

     80def eliminate_qualify(expression: exp.Expression) -> exp.Expression:
     81    """
     82    Convert SELECT statements that contain the QUALIFY clause into subqueries, filtered equivalently.
     83
     84    The idea behind this transformation can be seen in Snowflake's documentation for QUALIFY:
     85    https://docs.snowflake.com/en/sql-reference/constructs/qualify
     86
     87    Some dialects don't support window functions in the WHERE clause, so we need to include them as
     88    projections in the subquery, in order to refer to them in the outer filter using aliases. Also,
     89    if a column is referenced in the QUALIFY clause but is not selected, we need to include it too,
     90    otherwise we won't be able to refer to it in the outer query's WHERE clause.
     91    """
     92    if isinstance(expression, exp.Select) and expression.args.get("qualify"):
     93        taken = set(expression.named_selects)
     94        for select in expression.selects:
     95            if not select.alias_or_name:
     96                alias = find_new_name(taken, "_c")
     97                select.replace(exp.alias_(select.copy(), alias))
     98                taken.add(alias)
     99
    100        outer_selects = exp.select(*[select.alias_or_name for select in expression.selects])
    101        qualify_filters = expression.args["qualify"].pop().this
    102
    103        for expr in qualify_filters.find_all((exp.Window, exp.Column)):
    104            if isinstance(expr, exp.Window):
    105                alias = find_new_name(expression.named_selects, "_w")
    106                expression.select(exp.alias_(expr.copy(), alias), copy=False)
    107                column = exp.column(alias)
    108                if isinstance(expr.parent, exp.Qualify):
    109                    qualify_filters = column
    110                else:
    111                    expr.replace(column)
    112            elif expr.name not in expression.named_selects:
    113                expression.select(expr.copy(), copy=False)
    114
    115        return outer_selects.from_(expression.subquery(alias="_t")).where(qualify_filters)
    116
    117    return expression

Convert SELECT statements that contain the QUALIFY clause into subqueries, filtered equivalently.

The idea behind this transformation can be seen in Snowflake's documentation for QUALIFY: https://docs.snowflake.com/en/sql-reference/constructs/qualify

Some dialects don't support window functions in the WHERE clause, so we need to include them as projections in the subquery, in order to refer to them in the outer filter using aliases. Also, if a column is referenced in the QUALIFY clause but is not selected, we need to include it too, otherwise we won't be able to refer to it in the outer query's WHERE clause.

def remove_precision_parameterized_types( expression: sqlglot.expressions.Expression) -> sqlglot.expressions.Expression: View Source

    120def remove_precision_parameterized_types(expression: exp.Expression) -> exp.Expression:
    121    """
    122    Some dialects only allow the precision for parameterized types to be defined in the DDL and not in
    123    other expressions. This transforms removes the precision from parameterized types in expressions.
    124    """
    125    for node in expression.find_all(exp.DataType):
    126        node.set("expressions", [e for e in node.expressions if isinstance(e, exp.DataType)])
    127    return expression

Some dialects only allow the precision for parameterized types to be defined in the DDL and not in other expressions. This transforms removes the precision from parameterized types in expressions.

def unnest_to_explode( expression: sqlglot.expressions.Expression) -> sqlglot.expressions.Expression: View Source

    130def unnest_to_explode(expression: exp.Expression) -> exp.Expression:
    131    """Convert cross join unnest into lateral view explode (used in presto -> hive)."""
    132    if isinstance(expression, exp.Select):
    133        for join in expression.args.get("joins") or []:
    134            unnest = join.this
    135
    136            if isinstance(unnest, exp.Unnest):
    137                alias = unnest.args.get("alias")
    138                udtf = exp.Posexplode if unnest.args.get("ordinality") else exp.Explode
    139
    140                expression.args["joins"].remove(join)
    141
    142                for e, column in zip(unnest.expressions, alias.columns if alias else []):
    143                    expression.append(
    144                        "laterals",
    145                        exp.Lateral(
    146                            this=udtf(this=e),
    147                            view=True,
    148                            alias=exp.TableAlias(this=alias.this, columns=[column]),  # type: ignore
    149                        ),
    150                    )
    151    return expression

Convert cross join unnest into lateral view explode (used in presto -> hive).

def explode_to_unnest( expression: sqlglot.expressions.Expression) -> sqlglot.expressions.Expression: View Source

    154def explode_to_unnest(expression: exp.Expression) -> exp.Expression:
    155    """Convert explode/posexplode into unnest (used in hive -> presto)."""
    156    if isinstance(expression, exp.Select):
    157        from sqlglot.optimizer.scope import build_scope
    158
    159        taken_select_names = set(expression.named_selects)
    160        taken_source_names = set(build_scope(expression).selected_sources)
    161
    162        for select in expression.selects:
    163            to_replace = select
    164
    165            pos_alias = ""
    166            explode_alias = ""
    167
    168            if isinstance(select, exp.Alias):
    169                explode_alias = select.alias
    170                select = select.this
    171            elif isinstance(select, exp.Aliases):
    172                pos_alias = select.aliases[0].name
    173                explode_alias = select.aliases[1].name
    174                select = select.this
    175
    176            if isinstance(select, (exp.Explode, exp.Posexplode)):
    177                is_posexplode = isinstance(select, exp.Posexplode)
    178
    179                explode_arg = select.this
    180                unnest = exp.Unnest(expressions=[explode_arg.copy()], ordinality=is_posexplode)
    181
    182                # This ensures that we won't use [POS]EXPLODE's argument as a new selection
    183                if isinstance(explode_arg, exp.Column):
    184                    taken_select_names.add(explode_arg.output_name)
    185
    186                unnest_source_alias = find_new_name(taken_source_names, "_u")
    187                taken_source_names.add(unnest_source_alias)
    188
    189                if not explode_alias:
    190                    explode_alias = find_new_name(taken_select_names, "col")
    191                    taken_select_names.add(explode_alias)
    192
    193                    if is_posexplode:
    194                        pos_alias = find_new_name(taken_select_names, "pos")
    195                        taken_select_names.add(pos_alias)
    196
    197                if is_posexplode:
    198                    column_names = [explode_alias, pos_alias]
    199                    to_replace.pop()
    200                    expression.select(pos_alias, explode_alias, copy=False)
    201                else:
    202                    column_names = [explode_alias]
    203                    to_replace.replace(exp.column(explode_alias))
    204
    205                unnest = exp.alias_(unnest, unnest_source_alias, table=column_names)
    206
    207                if not expression.args.get("from"):
    208                    expression.from_(unnest, copy=False)
    209                else:
    210                    expression.join(unnest, join_type="CROSS", copy=False)
    211
    212    return expression

Convert explode/posexplode into unnest (used in hive -> presto).

def remove_target_from_merge( expression: sqlglot.expressions.Expression) -> sqlglot.expressions.Expression: View Source

    215def remove_target_from_merge(expression: exp.Expression) -> exp.Expression:
    216    """Remove table refs from columns in when statements."""
    217    if isinstance(expression, exp.Merge):
    218        alias = expression.this.args.get("alias")
    219        targets = {expression.this.this}
    220        if alias:
    221            targets.add(alias.this)
    222
    223        for when in expression.expressions:
    224            when.transform(
    225                lambda node: exp.column(node.name)
    226                if isinstance(node, exp.Column) and node.args.get("table") in targets
    227                else node,
    228                copy=False,
    229            )
    230    return expression

Remove table refs from columns in when statements.

def preprocess( transforms: List[Callable[[sqlglot.expressions.Expression], sqlglot.expressions.Expression]]) -> Callable[[sqlglot.generator.Generator, sqlglot.expressions.Expression], str]: View Source

    233def preprocess(
    234    transforms: t.List[t.Callable[[exp.Expression], exp.Expression]],
    235) -> t.Callable[[Generator, exp.Expression], str]:
    236    """
    237    Creates a new transform by chaining a sequence of transformations and converts the resulting
    238    expression to SQL, using an appropriate `Generator.TRANSFORMS` function.
    239
    240    Args:
    241        transforms: sequence of transform functions. These will be called in order.
    242
    243    Returns:
    244        Function that can be used as a generator transform.
    245    """
    246
    247    def _to_sql(self, expression: exp.Expression) -> str:
    248        expression = transforms[0](expression.copy())
    249        for t in transforms[1:]:
    250            expression = t(expression)
    251        return getattr(self, expression.key + "_sql")(expression)
    252
    253    return _to_sql

Creates a new transform by chaining a sequence of transformations and converts the resulting expression to SQL, using an appropriate `Generator.TRANSFORMS` function.

###### Arguments:

- transforms: sequence of transform functions. These will be called in order.

###### Returns:

> Function that can be used as a generator transform.