sqlglot.transforms API documentation

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    return expression.transform(
    125        lambda node: exp.DataType(
    126            **{
    127                **node.args,
    128                "expressions": [
    129                    node_expression
    130                    for node_expression in node.expressions
    131                    if isinstance(node_expression, exp.DataType)
    132                ],
    133            }
    134        )
    135        if isinstance(node, exp.DataType)
    136        else node,
    137    )
    138
    139
    140def unnest_to_explode(expression: exp.Expression) -> exp.Expression:
    141    """Convert cross join unnest into lateral view explode (used in presto -> hive)."""
    142    if isinstance(expression, exp.Select):
    143        for join in expression.args.get("joins") or []:
    144            unnest = join.this
    145
    146            if isinstance(unnest, exp.Unnest):
    147                alias = unnest.args.get("alias")
    148                udtf = exp.Posexplode if unnest.args.get("ordinality") else exp.Explode
    149
    150                expression.args["joins"].remove(join)
    151
    152                for e, column in zip(unnest.expressions, alias.columns if alias else []):
    153                    expression.append(
    154                        "laterals",
    155                        exp.Lateral(
    156                            this=udtf(this=e),
    157                            view=True,
    158                            alias=exp.TableAlias(this=alias.this, columns=[column]),  # type: ignore
    159                        ),
    160                    )
    161    return expression
    162
    163
    164def explode_to_unnest(expression: exp.Expression) -> exp.Expression:
    165    """Convert explode/posexplode into unnest (used in hive -> presto)."""
    166    if isinstance(expression, exp.Select):
    167        from sqlglot.optimizer.scope import build_scope
    168
    169        taken_select_names = set(expression.named_selects)
    170        taken_source_names = set(build_scope(expression).selected_sources)
    171
    172        for select in expression.selects:
    173            to_replace = select
    174
    175            pos_alias = ""
    176            explode_alias = ""
    177
    178            if isinstance(select, exp.Alias):
    179                explode_alias = select.alias
    180                select = select.this
    181            elif isinstance(select, exp.Aliases):
    182                pos_alias = select.aliases[0].name
    183                explode_alias = select.aliases[1].name
    184                select = select.this
    185
    186            if isinstance(select, (exp.Explode, exp.Posexplode)):
    187                is_posexplode = isinstance(select, exp.Posexplode)
    188
    189                explode_arg = select.this
    190                unnest = exp.Unnest(expressions=[explode_arg.copy()], ordinality=is_posexplode)
    191
    192                # This ensures that we won't use [POS]EXPLODE's argument as a new selection
    193                if isinstance(explode_arg, exp.Column):
    194                    taken_select_names.add(explode_arg.output_name)
    195
    196                unnest_source_alias = find_new_name(taken_source_names, "_u")
    197                taken_source_names.add(unnest_source_alias)
    198
    199                if not explode_alias:
    200                    explode_alias = find_new_name(taken_select_names, "col")
    201                    taken_select_names.add(explode_alias)
    202
    203                    if is_posexplode:
    204                        pos_alias = find_new_name(taken_select_names, "pos")
    205                        taken_select_names.add(pos_alias)
    206
    207                if is_posexplode:
    208                    column_names = [explode_alias, pos_alias]
    209                    to_replace.pop()
    210                    expression.select(pos_alias, explode_alias, copy=False)
    211                else:
    212                    column_names = [explode_alias]
    213                    to_replace.replace(exp.column(explode_alias))
    214
    215                unnest = exp.alias_(unnest, unnest_source_alias, table=column_names)
    216
    217                if not expression.args.get("from"):
    218                    expression.from_(unnest, copy=False)
    219                else:
    220                    expression.join(unnest, join_type="CROSS", copy=False)
    221
    222    return expression
    223
    224
    225def remove_target_from_merge(expression: exp.Expression) -> exp.Expression:
    226    """Remove table refs from columns in when statements."""
    227    if isinstance(expression, exp.Merge):
    228        alias = expression.this.args.get("alias")
    229        targets = {expression.this.this}
    230        if alias:
    231            targets.add(alias.this)
    232
    233        for when in expression.expressions:
    234            when.transform(
    235                lambda node: exp.column(node.name)
    236                if isinstance(node, exp.Column) and node.args.get("table") in targets
    237                else node,
    238                copy=False,
    239            )
    240    return expression
    241
    242
    243def preprocess(
    244    transforms: t.List[t.Callable[[exp.Expression], exp.Expression]],
    245) -> t.Callable[[Generator, exp.Expression], str]:
    246    """
    247    Creates a new transform by chaining a sequence of transformations and converts the resulting
    248    expression to SQL, using an appropriate `Generator.TRANSFORMS` function.
    249
    250    Args:
    251        transforms: sequence of transform functions. These will be called in order.
    252
    253    Returns:
    254        Function that can be used as a generator transform.
    255    """
    256
    257    def _to_sql(self, expression: exp.Expression) -> str:
    258        expression = transforms[0](expression.copy())
    259        for t in transforms[1:]:
    260            expression = t(expression)
    261        return getattr(self, expression.key + "_sql")(expression)
    262
    263    return _to_sql
    264
    265
    266UNALIAS_GROUP = {exp.Group: preprocess([unalias_group])}
    267ELIMINATE_DISTINCT_ON = {exp.Select: preprocess([eliminate_distinct_on])}
    268ELIMINATE_QUALIFY = {exp.Select: preprocess([eliminate_qualify])}
    269REMOVE_PRECISION_PARAMETERIZED_TYPES = {
    270    exp.Cast: preprocess([remove_precision_parameterized_types])
    271}

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    return expression.transform(
    126        lambda node: exp.DataType(
    127            **{
    128                **node.args,
    129                "expressions": [
    130                    node_expression
    131                    for node_expression in node.expressions
    132                    if isinstance(node_expression, exp.DataType)
    133                ],
    134            }
    135        )
    136        if isinstance(node, exp.DataType)
    137        else node,
    138    )

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

    141def unnest_to_explode(expression: exp.Expression) -> exp.Expression:
    142    """Convert cross join unnest into lateral view explode (used in presto -> hive)."""
    143    if isinstance(expression, exp.Select):
    144        for join in expression.args.get("joins") or []:
    145            unnest = join.this
    146
    147            if isinstance(unnest, exp.Unnest):
    148                alias = unnest.args.get("alias")
    149                udtf = exp.Posexplode if unnest.args.get("ordinality") else exp.Explode
    150
    151                expression.args["joins"].remove(join)
    152
    153                for e, column in zip(unnest.expressions, alias.columns if alias else []):
    154                    expression.append(
    155                        "laterals",
    156                        exp.Lateral(
    157                            this=udtf(this=e),
    158                            view=True,
    159                            alias=exp.TableAlias(this=alias.this, columns=[column]),  # type: ignore
    160                        ),
    161                    )
    162    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

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

    226def remove_target_from_merge(expression: exp.Expression) -> exp.Expression:
    227    """Remove table refs from columns in when statements."""
    228    if isinstance(expression, exp.Merge):
    229        alias = expression.this.args.get("alias")
    230        targets = {expression.this.this}
    231        if alias:
    232            targets.add(alias.this)
    233
    234        for when in expression.expressions:
    235            when.transform(
    236                lambda node: exp.column(node.name)
    237                if isinstance(node, exp.Column) and node.args.get("table") in targets
    238                else node,
    239                copy=False,
    240            )
    241    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

    244def preprocess(
    245    transforms: t.List[t.Callable[[exp.Expression], exp.Expression]],
    246) -> t.Callable[[Generator, exp.Expression], str]:
    247    """
    248    Creates a new transform by chaining a sequence of transformations and converts the resulting
    249    expression to SQL, using an appropriate `Generator.TRANSFORMS` function.
    250
    251    Args:
    252        transforms: sequence of transform functions. These will be called in order.
    253
    254    Returns:
    255        Function that can be used as a generator transform.
    256    """
    257
    258    def _to_sql(self, expression: exp.Expression) -> str:
    259        expression = transforms[0](expression.copy())
    260        for t in transforms[1:]:
    261            expression = t(expression)
    262        return getattr(self, expression.key + "_sql")(expression)
    263
    264    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.