sqlglot.transforms

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import expressions as exp
  6from sqlglot.helper import find_new_name, name_sequence
  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(this=exp.RowNumber(), partition_by=distinct_cols)
 67        order = expression.args.get("order")
 68
 69        if order:
 70            window.set("order", order.pop())
 71        else:
 72            window.set("order", exp.Order(expressions=[c.copy() for c in distinct_cols]))
 73
 74        window = exp.alias_(window, row_number)
 75        expression.select(window, copy=False)
 76
 77        return (
 78            exp.select(*outer_selects, copy=False)
 79            .from_(expression.subquery("_t", copy=False), copy=False)
 80            .where(exp.column(row_number).eq(1), copy=False)
 81        )
 82
 83    return expression
 84
 85
 86def eliminate_qualify(expression: exp.Expression) -> exp.Expression:
 87    """
 88    Convert SELECT statements that contain the QUALIFY clause into subqueries, filtered equivalently.
 89
 90    The idea behind this transformation can be seen in Snowflake's documentation for QUALIFY:
 91    https://docs.snowflake.com/en/sql-reference/constructs/qualify
 92
 93    Some dialects don't support window functions in the WHERE clause, so we need to include them as
 94    projections in the subquery, in order to refer to them in the outer filter using aliases. Also,
 95    if a column is referenced in the QUALIFY clause but is not selected, we need to include it too,
 96    otherwise we won't be able to refer to it in the outer query's WHERE clause. Finally, if a
 97    newly aliased projection is referenced in the QUALIFY clause, it will be replaced by the
 98    corresponding expression to avoid creating invalid column references.
 99    """
100    if isinstance(expression, exp.Select) and expression.args.get("qualify"):
101        taken = set(expression.named_selects)
102        for select in expression.selects:
103            if not select.alias_or_name:
104                alias = find_new_name(taken, "_c")
105                select.replace(exp.alias_(select, alias))
106                taken.add(alias)
107
108        outer_selects = exp.select(*[select.alias_or_name for select in expression.selects])
109        qualify_filters = expression.args["qualify"].pop().this
110        expression_by_alias = {
111            select.alias: select.this
112            for select in expression.selects
113            if isinstance(select, exp.Alias)
114        }
115
116        select_candidates = exp.Window if expression.is_star else (exp.Window, exp.Column)
117        for select_candidate in qualify_filters.find_all(select_candidates):
118            if isinstance(select_candidate, exp.Window):
119                if expression_by_alias:
120                    for column in select_candidate.find_all(exp.Column):
121                        expr = expression_by_alias.get(column.name)
122                        if expr:
123                            column.replace(expr)
124
125                alias = find_new_name(expression.named_selects, "_w")
126                expression.select(exp.alias_(select_candidate, alias), copy=False)
127                column = exp.column(alias)
128
129                if isinstance(select_candidate.parent, exp.Qualify):
130                    qualify_filters = column
131                else:
132                    select_candidate.replace(column)
133            elif select_candidate.name not in expression.named_selects:
134                expression.select(select_candidate.copy(), copy=False)
135
136        return outer_selects.from_(expression.subquery(alias="_t", copy=False), copy=False).where(
137            qualify_filters, copy=False
138        )
139
140    return expression
141
142
143def remove_precision_parameterized_types(expression: exp.Expression) -> exp.Expression:
144    """
145    Some dialects only allow the precision for parameterized types to be defined in the DDL and not in
146    other expressions. This transforms removes the precision from parameterized types in expressions.
147    """
148    for node in expression.find_all(exp.DataType):
149        node.set(
150            "expressions", [e for e in node.expressions if not isinstance(e, exp.DataTypeParam)]
151        )
152
153    return expression
154
155
156def unqualify_unnest(expression: exp.Expression) -> exp.Expression:
157    """Remove references to unnest table aliases, added by the optimizer's qualify_columns step."""
158    from sqlglot.optimizer.scope import find_all_in_scope
159
160    if isinstance(expression, exp.Select):
161        unnest_aliases = {
162            unnest.alias
163            for unnest in find_all_in_scope(expression, exp.Unnest)
164            if isinstance(unnest.parent, (exp.From, exp.Join))
165        }
166        if unnest_aliases:
167            for column in expression.find_all(exp.Column):
168                if column.table in unnest_aliases:
169                    column.set("table", None)
170                elif column.db in unnest_aliases:
171                    column.set("db", None)
172
173    return expression
174
175
176def unnest_to_explode(expression: exp.Expression) -> exp.Expression:
177    """Convert cross join unnest into lateral view explode."""
178    if isinstance(expression, exp.Select):
179        for join in expression.args.get("joins") or []:
180            unnest = join.this
181
182            if isinstance(unnest, exp.Unnest):
183                alias = unnest.args.get("alias")
184                udtf = exp.Posexplode if unnest.args.get("offset") else exp.Explode
185
186                expression.args["joins"].remove(join)
187
188                for e, column in zip(unnest.expressions, alias.columns if alias else []):
189                    expression.append(
190                        "laterals",
191                        exp.Lateral(
192                            this=udtf(this=e),
193                            view=True,
194                            alias=exp.TableAlias(this=alias.this, columns=[column]),  # type: ignore
195                        ),
196                    )
197
198    return expression
199
200
201def explode_to_unnest(index_offset: int = 0) -> t.Callable[[exp.Expression], exp.Expression]:
202    """Convert explode/posexplode into unnest."""
203
204    def _explode_to_unnest(expression: exp.Expression) -> exp.Expression:
205        if isinstance(expression, exp.Select):
206            from sqlglot.optimizer.scope import Scope
207
208            taken_select_names = set(expression.named_selects)
209            taken_source_names = {name for name, _ in Scope(expression).references}
210
211            def new_name(names: t.Set[str], name: str) -> str:
212                name = find_new_name(names, name)
213                names.add(name)
214                return name
215
216            arrays: t.List[exp.Condition] = []
217            series_alias = new_name(taken_select_names, "pos")
218            series = exp.alias_(
219                exp.Unnest(
220                    expressions=[exp.GenerateSeries(start=exp.Literal.number(index_offset))]
221                ),
222                new_name(taken_source_names, "_u"),
223                table=[series_alias],
224            )
225
226            # we use list here because expression.selects is mutated inside the loop
227            for select in list(expression.selects):
228                explode = select.find(exp.Explode)
229
230                if explode:
231                    pos_alias = ""
232                    explode_alias = ""
233
234                    if isinstance(select, exp.Alias):
235                        explode_alias = select.args["alias"]
236                        alias = select
237                    elif isinstance(select, exp.Aliases):
238                        pos_alias = select.aliases[0]
239                        explode_alias = select.aliases[1]
240                        alias = select.replace(exp.alias_(select.this, "", copy=False))
241                    else:
242                        alias = select.replace(exp.alias_(select, ""))
243                        explode = alias.find(exp.Explode)
244                        assert explode
245
246                    is_posexplode = isinstance(explode, exp.Posexplode)
247                    explode_arg = explode.this
248
249                    if isinstance(explode, exp.ExplodeOuter):
250                        bracket = explode_arg[0]
251                        bracket.set("safe", True)
252                        bracket.set("offset", True)
253                        explode_arg = exp.func(
254                            "IF",
255                            exp.func(
256                                "ARRAY_SIZE", exp.func("COALESCE", explode_arg, exp.Array())
257                            ).eq(0),
258                            exp.array(bracket, copy=False),
259                            explode_arg,
260                        )
261
262                    # This ensures that we won't use [POS]EXPLODE's argument as a new selection
263                    if isinstance(explode_arg, exp.Column):
264                        taken_select_names.add(explode_arg.output_name)
265
266                    unnest_source_alias = new_name(taken_source_names, "_u")
267
268                    if not explode_alias:
269                        explode_alias = new_name(taken_select_names, "col")
270
271                        if is_posexplode:
272                            pos_alias = new_name(taken_select_names, "pos")
273
274                    if not pos_alias:
275                        pos_alias = new_name(taken_select_names, "pos")
276
277                    alias.set("alias", exp.to_identifier(explode_alias))
278
279                    series_table_alias = series.args["alias"].this
280                    column = exp.If(
281                        this=exp.column(series_alias, table=series_table_alias).eq(
282                            exp.column(pos_alias, table=unnest_source_alias)
283                        ),
284                        true=exp.column(explode_alias, table=unnest_source_alias),
285                    )
286
287                    explode.replace(column)
288
289                    if is_posexplode:
290                        expressions = expression.expressions
291                        expressions.insert(
292                            expressions.index(alias) + 1,
293                            exp.If(
294                                this=exp.column(series_alias, table=series_table_alias).eq(
295                                    exp.column(pos_alias, table=unnest_source_alias)
296                                ),
297                                true=exp.column(pos_alias, table=unnest_source_alias),
298                            ).as_(pos_alias),
299                        )
300                        expression.set("expressions", expressions)
301
302                    if not arrays:
303                        if expression.args.get("from"):
304                            expression.join(series, copy=False, join_type="CROSS")
305                        else:
306                            expression.from_(series, copy=False)
307
308                    size: exp.Condition = exp.ArraySize(this=explode_arg.copy())
309                    arrays.append(size)
310
311                    # trino doesn't support left join unnest with on conditions
312                    # if it did, this would be much simpler
313                    expression.join(
314                        exp.alias_(
315                            exp.Unnest(
316                                expressions=[explode_arg.copy()],
317                                offset=exp.to_identifier(pos_alias),
318                            ),
319                            unnest_source_alias,
320                            table=[explode_alias],
321                        ),
322                        join_type="CROSS",
323                        copy=False,
324                    )
325
326                    if index_offset != 1:
327                        size = size - 1
328
329                    expression.where(
330                        exp.column(series_alias, table=series_table_alias)
331                        .eq(exp.column(pos_alias, table=unnest_source_alias))
332                        .or_(
333                            (exp.column(series_alias, table=series_table_alias) > size).and_(
334                                exp.column(pos_alias, table=unnest_source_alias).eq(size)
335                            )
336                        ),
337                        copy=False,
338                    )
339
340            if arrays:
341                end: exp.Condition = exp.Greatest(this=arrays[0], expressions=arrays[1:])
342
343                if index_offset != 1:
344                    end = end - (1 - index_offset)
345                series.expressions[0].set("end", end)
346
347        return expression
348
349    return _explode_to_unnest
350
351
352def add_within_group_for_percentiles(expression: exp.Expression) -> exp.Expression:
353    """Transforms percentiles by adding a WITHIN GROUP clause to them."""
354    if (
355        isinstance(expression, exp.PERCENTILES)
356        and not isinstance(expression.parent, exp.WithinGroup)
357        and expression.expression
358    ):
359        column = expression.this.pop()
360        expression.set("this", expression.expression.pop())
361        order = exp.Order(expressions=[exp.Ordered(this=column)])
362        expression = exp.WithinGroup(this=expression, expression=order)
363
364    return expression
365
366
367def remove_within_group_for_percentiles(expression: exp.Expression) -> exp.Expression:
368    """Transforms percentiles by getting rid of their corresponding WITHIN GROUP clause."""
369    if (
370        isinstance(expression, exp.WithinGroup)
371        and isinstance(expression.this, exp.PERCENTILES)
372        and isinstance(expression.expression, exp.Order)
373    ):
374        quantile = expression.this.this
375        input_value = t.cast(exp.Ordered, expression.find(exp.Ordered)).this
376        return expression.replace(exp.ApproxQuantile(this=input_value, quantile=quantile))
377
378    return expression
379
380
381def add_recursive_cte_column_names(expression: exp.Expression) -> exp.Expression:
382    """Uses projection output names in recursive CTE definitions to define the CTEs' columns."""
383    if isinstance(expression, exp.With) and expression.recursive:
384        next_name = name_sequence("_c_")
385
386        for cte in expression.expressions:
387            if not cte.args["alias"].columns:
388                query = cte.this
389                if isinstance(query, exp.Union):
390                    query = query.this
391
392                cte.args["alias"].set(
393                    "columns",
394                    [exp.to_identifier(s.alias_or_name or next_name()) for s in query.selects],
395                )
396
397    return expression
398
399
400def epoch_cast_to_ts(expression: exp.Expression) -> exp.Expression:
401    """Replace 'epoch' in casts by the equivalent date literal."""
402    if (
403        isinstance(expression, (exp.Cast, exp.TryCast))
404        and expression.name.lower() == "epoch"
405        and expression.to.this in exp.DataType.TEMPORAL_TYPES
406    ):
407        expression.this.replace(exp.Literal.string("1970-01-01 00:00:00"))
408
409    return expression
410
411
412def eliminate_semi_and_anti_joins(expression: exp.Expression) -> exp.Expression:
413    """Convert SEMI and ANTI joins into equivalent forms that use EXIST instead."""
414    if isinstance(expression, exp.Select):
415        for join in expression.args.get("joins") or []:
416            on = join.args.get("on")
417            if on and join.kind in ("SEMI", "ANTI"):
418                subquery = exp.select("1").from_(join.this).where(on)
419                exists = exp.Exists(this=subquery)
420                if join.kind == "ANTI":
421                    exists = exists.not_(copy=False)
422
423                join.pop()
424                expression.where(exists, copy=False)
425
426    return expression
427
428
429def eliminate_full_outer_join(expression: exp.Expression) -> exp.Expression:
430    """
431    Converts a query with a FULL OUTER join to a union of identical queries that
432    use LEFT/RIGHT OUTER joins instead. This transformation currently only works
433    for queries that have a single FULL OUTER join.
434    """
435    if isinstance(expression, exp.Select):
436        full_outer_joins = [
437            (index, join)
438            for index, join in enumerate(expression.args.get("joins") or [])
439            if join.side == "FULL"
440        ]
441
442        if len(full_outer_joins) == 1:
443            expression_copy = expression.copy()
444            expression.set("limit", None)
445            index, full_outer_join = full_outer_joins[0]
446            full_outer_join.set("side", "left")
447            expression_copy.args["joins"][index].set("side", "right")
448            expression_copy.args.pop("with", None)  # remove CTEs from RIGHT side
449
450            return exp.union(expression, expression_copy, copy=False)
451
452    return expression
453
454
455def move_ctes_to_top_level(expression: exp.Expression) -> exp.Expression:
456    """
457    Some dialects (e.g. Hive, T-SQL, Spark prior to version 3) only allow CTEs to be
458    defined at the top-level, so for example queries like:
459
460        SELECT * FROM (WITH t(c) AS (SELECT 1) SELECT * FROM t) AS subq
461
462    are invalid in those dialects. This transformation can be used to ensure all CTEs are
463    moved to the top level so that the final SQL code is valid from a syntax standpoint.
464
465    TODO: handle name clashes whilst moving CTEs (it can get quite tricky & costly).
466    """
467    top_level_with = expression.args.get("with")
468    for node in expression.find_all(exp.With):
469        if node.parent is expression:
470            continue
471
472        inner_with = node.pop()
473        if not top_level_with:
474            top_level_with = inner_with
475            expression.set("with", top_level_with)
476        else:
477            if inner_with.recursive:
478                top_level_with.set("recursive", True)
479
480            top_level_with.set("expressions", inner_with.expressions + top_level_with.expressions)
481
482    return expression
483
484
485def ensure_bools(expression: exp.Expression) -> exp.Expression:
486    """Converts numeric values used in conditions into explicit boolean expressions."""
487    from sqlglot.optimizer.canonicalize import ensure_bools
488
489    def _ensure_bool(node: exp.Expression) -> None:
490        if (
491            node.is_number
492            or node.is_type(exp.DataType.Type.UNKNOWN, *exp.DataType.NUMERIC_TYPES)
493            or (isinstance(node, exp.Column) and not node.type)
494        ):
495            node.replace(node.neq(0))
496
497    for node in expression.walk():
498        ensure_bools(node, _ensure_bool)
499
500    return expression
501
502
503def unqualify_columns(expression: exp.Expression) -> exp.Expression:
504    for column in expression.find_all(exp.Column):
505        # We only wanna pop off the table, db, catalog args
506        for part in column.parts[:-1]:
507            part.pop()
508
509    return expression
510
511
512def remove_unique_constraints(expression: exp.Expression) -> exp.Expression:
513    assert isinstance(expression, exp.Create)
514    for constraint in expression.find_all(exp.UniqueColumnConstraint):
515        if constraint.parent:
516            constraint.parent.pop()
517
518    return expression
519
520
521def ctas_with_tmp_tables_to_create_tmp_view(
522    expression: exp.Expression,
523    tmp_storage_provider: t.Callable[[exp.Expression], exp.Expression] = lambda e: e,
524) -> exp.Expression:
525    assert isinstance(expression, exp.Create)
526    properties = expression.args.get("properties")
527    temporary = any(
528        isinstance(prop, exp.TemporaryProperty)
529        for prop in (properties.expressions if properties else [])
530    )
531
532    # CTAS with temp tables map to CREATE TEMPORARY VIEW
533    if expression.kind == "TABLE" and temporary:
534        if expression.expression:
535            return exp.Create(
536                kind="TEMPORARY VIEW",
537                this=expression.this,
538                expression=expression.expression,
539            )
540        return tmp_storage_provider(expression)
541
542    return expression
543
544
545def move_schema_columns_to_partitioned_by(expression: exp.Expression) -> exp.Expression:
546    """
547    In Hive, the PARTITIONED BY property acts as an extension of a table's schema. When the
548    PARTITIONED BY value is an array of column names, they are transformed into a schema.
549    The corresponding columns are removed from the create statement.
550    """
551    assert isinstance(expression, exp.Create)
552    has_schema = isinstance(expression.this, exp.Schema)
553    is_partitionable = expression.kind in {"TABLE", "VIEW"}
554
555    if has_schema and is_partitionable:
556        prop = expression.find(exp.PartitionedByProperty)
557        if prop and prop.this and not isinstance(prop.this, exp.Schema):
558            schema = expression.this
559            columns = {v.name.upper() for v in prop.this.expressions}
560            partitions = [col for col in schema.expressions if col.name.upper() in columns]
561            schema.set("expressions", [e for e in schema.expressions if e not in partitions])
562            prop.replace(exp.PartitionedByProperty(this=exp.Schema(expressions=partitions)))
563            expression.set("this", schema)
564
565    return expression
566
567
568def move_partitioned_by_to_schema_columns(expression: exp.Expression) -> exp.Expression:
569    """
570    Spark 3 supports both "HIVEFORMAT" and "DATASOURCE" formats for CREATE TABLE.
571
572    Currently, SQLGlot uses the DATASOURCE format for Spark 3.
573    """
574    assert isinstance(expression, exp.Create)
575    prop = expression.find(exp.PartitionedByProperty)
576    if (
577        prop
578        and prop.this
579        and isinstance(prop.this, exp.Schema)
580        and all(isinstance(e, exp.ColumnDef) and e.kind for e in prop.this.expressions)
581    ):
582        prop_this = exp.Tuple(
583            expressions=[exp.to_identifier(e.this) for e in prop.this.expressions]
584        )
585        schema = expression.this
586        for e in prop.this.expressions:
587            schema.append("expressions", e)
588        prop.set("this", prop_this)
589
590    return expression
591
592
593def struct_kv_to_alias(expression: exp.Expression) -> exp.Expression:
594    """Converts struct arguments to aliases, e.g. STRUCT(1 AS y)."""
595    if isinstance(expression, exp.Struct):
596        expression.set(
597            "expressions",
598            [
599                exp.alias_(e.expression, e.this) if isinstance(e, exp.PropertyEQ) else e
600                for e in expression.expressions
601            ],
602        )
603
604    return expression
605
606
607def preprocess(
608    transforms: t.List[t.Callable[[exp.Expression], exp.Expression]],
609) -> t.Callable[[Generator, exp.Expression], str]:
610    """
611    Creates a new transform by chaining a sequence of transformations and converts the resulting
612    expression to SQL, using either the "_sql" method corresponding to the resulting expression,
613    or the appropriate `Generator.TRANSFORMS` function (when applicable -- see below).
614
615    Args:
616        transforms: sequence of transform functions. These will be called in order.
617
618    Returns:
619        Function that can be used as a generator transform.
620    """
621
622    def _to_sql(self, expression: exp.Expression) -> str:
623        expression_type = type(expression)
624
625        expression = transforms[0](expression)
626        for transform in transforms[1:]:
627            expression = transform(expression)
628
629        _sql_handler = getattr(self, expression.key + "_sql", None)
630        if _sql_handler:
631            return _sql_handler(expression)
632
633        transforms_handler = self.TRANSFORMS.get(type(expression))
634        if transforms_handler:
635            if expression_type is type(expression):
636                if isinstance(expression, exp.Func):
637                    return self.function_fallback_sql(expression)
638
639                # Ensures we don't enter an infinite loop. This can happen when the original expression
640                # has the same type as the final expression and there's no _sql method available for it,
641                # because then it'd re-enter _to_sql.
642                raise ValueError(
643                    f"Expression type {expression.__class__.__name__} requires a _sql method in order to be transformed."
644                )
645
646            return transforms_handler(self, expression)
647
648        raise ValueError(f"Unsupported expression type {expression.__class__.__name__}.")
649
650    return _to_sql