sqlglot.transforms

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import expressions as exp
  6from sqlglot.errors import UnsupportedError
  7from sqlglot.helper import find_new_name, name_sequence
  8
  9
 10if t.TYPE_CHECKING:
 11    from sqlglot._typing import E
 12    from sqlglot.generator import Generator
 13
 14
 15def preprocess(
 16    transforms: t.List[t.Callable[[exp.Expression], exp.Expression]],
 17) -> t.Callable[[Generator, exp.Expression], str]:
 18    """
 19    Creates a new transform by chaining a sequence of transformations and converts the resulting
 20    expression to SQL, using either the "_sql" method corresponding to the resulting expression,
 21    or the appropriate `Generator.TRANSFORMS` function (when applicable -- see below).
 22
 23    Args:
 24        transforms: sequence of transform functions. These will be called in order.
 25
 26    Returns:
 27        Function that can be used as a generator transform.
 28    """
 29
 30    def _to_sql(self, expression: exp.Expression) -> str:
 31        expression_type = type(expression)
 32
 33        try:
 34            expression = transforms[0](expression)
 35            for transform in transforms[1:]:
 36                expression = transform(expression)
 37        except UnsupportedError as unsupported_error:
 38            self.unsupported(str(unsupported_error))
 39
 40        _sql_handler = getattr(self, expression.key + "_sql", None)
 41        if _sql_handler:
 42            return _sql_handler(expression)
 43
 44        transforms_handler = self.TRANSFORMS.get(type(expression))
 45        if transforms_handler:
 46            if expression_type is type(expression):
 47                if isinstance(expression, exp.Func):
 48                    return self.function_fallback_sql(expression)
 49
 50                # Ensures we don't enter an infinite loop. This can happen when the original expression
 51                # has the same type as the final expression and there's no _sql method available for it,
 52                # because then it'd re-enter _to_sql.
 53                raise ValueError(
 54                    f"Expression type {expression.__class__.__name__} requires a _sql method in order to be transformed."
 55                )
 56
 57            return transforms_handler(self, expression)
 58
 59        raise ValueError(f"Unsupported expression type {expression.__class__.__name__}.")
 60
 61    return _to_sql
 62
 63
 64def unnest_generate_date_array_using_recursive_cte(expression: exp.Expression) -> exp.Expression:
 65    if isinstance(expression, exp.Select):
 66        count = 0
 67        recursive_ctes = []
 68
 69        for unnest in expression.find_all(exp.Unnest):
 70            if (
 71                not isinstance(unnest.parent, (exp.From, exp.Join))
 72                or len(unnest.expressions) != 1
 73                or not isinstance(unnest.expressions[0], exp.GenerateDateArray)
 74            ):
 75                continue
 76
 77            generate_date_array = unnest.expressions[0]
 78            start = generate_date_array.args.get("start")
 79            end = generate_date_array.args.get("end")
 80            step = generate_date_array.args.get("step")
 81
 82            if not start or not end or not isinstance(step, exp.Interval):
 83                continue
 84
 85            alias = unnest.args.get("alias")
 86            column_name = alias.columns[0] if isinstance(alias, exp.TableAlias) else "date_value"
 87
 88            start = exp.cast(start, "date")
 89            date_add = exp.func(
 90                "date_add", column_name, exp.Literal.number(step.name), step.args.get("unit")
 91            )
 92            cast_date_add = exp.cast(date_add, "date")
 93
 94            cte_name = "_generated_dates" + (f"_{count}" if count else "")
 95
 96            base_query = exp.select(start.as_(column_name))
 97            recursive_query = (
 98                exp.select(cast_date_add)
 99                .from_(cte_name)
100                .where(cast_date_add <= exp.cast(end, "date"))
101            )
102            cte_query = base_query.union(recursive_query, distinct=False)
103
104            generate_dates_query = exp.select(column_name).from_(cte_name)
105            unnest.replace(generate_dates_query.subquery(cte_name))
106
107            recursive_ctes.append(
108                exp.alias_(exp.CTE(this=cte_query), cte_name, table=[column_name])
109            )
110            count += 1
111
112        if recursive_ctes:
113            with_expression = expression.args.get("with") or exp.With()
114            with_expression.set("recursive", True)
115            with_expression.set("expressions", [*recursive_ctes, *with_expression.expressions])
116            expression.set("with", with_expression)
117
118    return expression
119
120
121def unnest_generate_series(expression: exp.Expression) -> exp.Expression:
122    """Unnests GENERATE_SERIES or SEQUENCE table references."""
123    this = expression.this
124    if isinstance(expression, exp.Table) and isinstance(this, exp.GenerateSeries):
125        unnest = exp.Unnest(expressions=[this])
126        if expression.alias:
127            return exp.alias_(unnest, alias="_u", table=[expression.alias], copy=False)
128
129        return unnest
130
131    return expression
132
133
134def unalias_group(expression: exp.Expression) -> exp.Expression:
135    """
136    Replace references to select aliases in GROUP BY clauses.
137
138    Example:
139        >>> import sqlglot
140        >>> sqlglot.parse_one("SELECT a AS b FROM x GROUP BY b").transform(unalias_group).sql()
141        'SELECT a AS b FROM x GROUP BY 1'
142
143    Args:
144        expression: the expression that will be transformed.
145
146    Returns:
147        The transformed expression.
148    """
149    if isinstance(expression, exp.Group) and isinstance(expression.parent, exp.Select):
150        aliased_selects = {
151            e.alias: i
152            for i, e in enumerate(expression.parent.expressions, start=1)
153            if isinstance(e, exp.Alias)
154        }
155
156        for group_by in expression.expressions:
157            if (
158                isinstance(group_by, exp.Column)
159                and not group_by.table
160                and group_by.name in aliased_selects
161            ):
162                group_by.replace(exp.Literal.number(aliased_selects.get(group_by.name)))
163
164    return expression
165
166
167def eliminate_distinct_on(expression: exp.Expression) -> exp.Expression:
168    """
169    Convert SELECT DISTINCT ON statements to a subquery with a window function.
170
171    This is useful for dialects that don't support SELECT DISTINCT ON but support window functions.
172
173    Args:
174        expression: the expression that will be transformed.
175
176    Returns:
177        The transformed expression.
178    """
179    if (
180        isinstance(expression, exp.Select)
181        and expression.args.get("distinct")
182        and isinstance(expression.args["distinct"].args.get("on"), exp.Tuple)
183    ):
184        row_number_window_alias = find_new_name(expression.named_selects, "_row_number")
185
186        distinct_cols = expression.args["distinct"].pop().args["on"].expressions
187        window = exp.Window(this=exp.RowNumber(), partition_by=distinct_cols)
188
189        order = expression.args.get("order")
190        if order:
191            window.set("order", order.pop())
192        else:
193            window.set("order", exp.Order(expressions=[c.copy() for c in distinct_cols]))
194
195        window = exp.alias_(window, row_number_window_alias)
196        expression.select(window, copy=False)
197
198        # We add aliases to the projections so that we can safely reference them in the outer query
199        new_selects = []
200        taken_names = {row_number_window_alias}
201        for select in expression.selects[:-1]:
202            if select.is_star:
203                new_selects = [exp.Star()]
204                break
205
206            if not isinstance(select, exp.Alias):
207                alias = find_new_name(taken_names, select.output_name or "_col")
208                quoted = select.this.args.get("quoted") if isinstance(select, exp.Column) else None
209                select = select.replace(exp.alias_(select, alias, quoted=quoted))
210
211            taken_names.add(select.output_name)
212            new_selects.append(select.args["alias"])
213
214        return (
215            exp.select(*new_selects, copy=False)
216            .from_(expression.subquery("_t", copy=False), copy=False)
217            .where(exp.column(row_number_window_alias).eq(1), copy=False)
218        )
219
220    return expression
221
222
223def eliminate_qualify(expression: exp.Expression) -> exp.Expression:
224    """
225    Convert SELECT statements that contain the QUALIFY clause into subqueries, filtered equivalently.
226
227    The idea behind this transformation can be seen in Snowflake's documentation for QUALIFY:
228    https://docs.snowflake.com/en/sql-reference/constructs/qualify
229
230    Some dialects don't support window functions in the WHERE clause, so we need to include them as
231    projections in the subquery, in order to refer to them in the outer filter using aliases. Also,
232    if a column is referenced in the QUALIFY clause but is not selected, we need to include it too,
233    otherwise we won't be able to refer to it in the outer query's WHERE clause. Finally, if a
234    newly aliased projection is referenced in the QUALIFY clause, it will be replaced by the
235    corresponding expression to avoid creating invalid column references.
236    """
237    if isinstance(expression, exp.Select) and expression.args.get("qualify"):
238        taken = set(expression.named_selects)
239        for select in expression.selects:
240            if not select.alias_or_name:
241                alias = find_new_name(taken, "_c")
242                select.replace(exp.alias_(select, alias))
243                taken.add(alias)
244
245        def _select_alias_or_name(select: exp.Expression) -> str | exp.Column:
246            alias_or_name = select.alias_or_name
247            identifier = select.args.get("alias") or select.this
248            if isinstance(identifier, exp.Identifier):
249                return exp.column(alias_or_name, quoted=identifier.args.get("quoted"))
250            return alias_or_name
251
252        outer_selects = exp.select(*list(map(_select_alias_or_name, expression.selects)))
253        qualify_filters = expression.args["qualify"].pop().this
254        expression_by_alias = {
255            select.alias: select.this
256            for select in expression.selects
257            if isinstance(select, exp.Alias)
258        }
259
260        select_candidates = exp.Window if expression.is_star else (exp.Window, exp.Column)
261        for select_candidate in list(qualify_filters.find_all(select_candidates)):
262            if isinstance(select_candidate, exp.Window):
263                if expression_by_alias:
264                    for column in select_candidate.find_all(exp.Column):
265                        expr = expression_by_alias.get(column.name)
266                        if expr:
267                            column.replace(expr)
268
269                alias = find_new_name(expression.named_selects, "_w")
270                expression.select(exp.alias_(select_candidate, alias), copy=False)
271                column = exp.column(alias)
272
273                if isinstance(select_candidate.parent, exp.Qualify):
274                    qualify_filters = column
275                else:
276                    select_candidate.replace(column)
277            elif select_candidate.name not in expression.named_selects:
278                expression.select(select_candidate.copy(), copy=False)
279
280        return outer_selects.from_(expression.subquery(alias="_t", copy=False), copy=False).where(
281            qualify_filters, copy=False
282        )
283
284    return expression
285
286
287def remove_precision_parameterized_types(expression: exp.Expression) -> exp.Expression:
288    """
289    Some dialects only allow the precision for parameterized types to be defined in the DDL and not in
290    other expressions. This transforms removes the precision from parameterized types in expressions.
291    """
292    for node in expression.find_all(exp.DataType):
293        node.set(
294            "expressions", [e for e in node.expressions if not isinstance(e, exp.DataTypeParam)]
295        )
296
297    return expression
298
299
300def unqualify_unnest(expression: exp.Expression) -> exp.Expression:
301    """Remove references to unnest table aliases, added by the optimizer's qualify_columns step."""
302    from sqlglot.optimizer.scope import find_all_in_scope
303
304    if isinstance(expression, exp.Select):
305        unnest_aliases = {
306            unnest.alias
307            for unnest in find_all_in_scope(expression, exp.Unnest)
308            if isinstance(unnest.parent, (exp.From, exp.Join))
309        }
310        if unnest_aliases:
311            for column in expression.find_all(exp.Column):
312                if column.table in unnest_aliases:
313                    column.set("table", None)
314                elif column.db in unnest_aliases:
315                    column.set("db", None)
316
317    return expression
318
319
320def unnest_to_explode(
321    expression: exp.Expression,
322    unnest_using_arrays_zip: bool = True,
323) -> exp.Expression:
324    """Convert cross join unnest into lateral view explode."""
325
326    def _unnest_zip_exprs(
327        u: exp.Unnest, unnest_exprs: t.List[exp.Expression], has_multi_expr: bool
328    ) -> t.List[exp.Expression]:
329        if has_multi_expr:
330            if not unnest_using_arrays_zip:
331                raise UnsupportedError("Cannot transpile UNNEST with multiple input arrays")
332
333            # Use INLINE(ARRAYS_ZIP(...)) for multiple expressions
334            zip_exprs: t.List[exp.Expression] = [
335                exp.Anonymous(this="ARRAYS_ZIP", expressions=unnest_exprs)
336            ]
337            u.set("expressions", zip_exprs)
338            return zip_exprs
339        return unnest_exprs
340
341    def _udtf_type(u: exp.Unnest, has_multi_expr: bool) -> t.Type[exp.Func]:
342        if u.args.get("offset"):
343            return exp.Posexplode
344        return exp.Inline if has_multi_expr else exp.Explode
345
346    if isinstance(expression, exp.Select):
347        from_ = expression.args.get("from")
348
349        if from_ and isinstance(from_.this, exp.Unnest):
350            unnest = from_.this
351            alias = unnest.args.get("alias")
352            exprs = unnest.expressions
353            has_multi_expr = len(exprs) > 1
354            this, *expressions = _unnest_zip_exprs(unnest, exprs, has_multi_expr)
355
356            unnest.replace(
357                exp.Table(
358                    this=_udtf_type(unnest, has_multi_expr)(
359                        this=this,
360                        expressions=expressions,
361                    ),
362                    alias=exp.TableAlias(this=alias.this, columns=alias.columns) if alias else None,
363                )
364            )
365
366        joins = expression.args.get("joins") or []
367        for join in list(joins):
368            join_expr = join.this
369
370            is_lateral = isinstance(join_expr, exp.Lateral)
371
372            unnest = join_expr.this if is_lateral else join_expr
373
374            if isinstance(unnest, exp.Unnest):
375                if is_lateral:
376                    alias = join_expr.args.get("alias")
377                else:
378                    alias = unnest.args.get("alias")
379                exprs = unnest.expressions
380                # The number of unnest.expressions will be changed by _unnest_zip_exprs, we need to record it here
381                has_multi_expr = len(exprs) > 1
382                exprs = _unnest_zip_exprs(unnest, exprs, has_multi_expr)
383
384                joins.remove(join)
385
386                alias_cols = alias.columns if alias else []
387
388                # # Handle UNNEST to LATERAL VIEW EXPLODE: Exception is raised when there are 0 or > 2 aliases
389                # Spark LATERAL VIEW EXPLODE requires single alias for array/struct and two for Map type column unlike unnest in trino/presto which can take an arbitrary amount.
390                # Refs: https://spark.apache.org/docs/latest/sql-ref-syntax-qry-select-lateral-view.html
391
392                if not has_multi_expr and len(alias_cols) not in (1, 2):
393                    raise UnsupportedError(
394                        "CROSS JOIN UNNEST to LATERAL VIEW EXPLODE transformation requires explicit column aliases"
395                    )
396
397                for e, column in zip(exprs, alias_cols):
398                    expression.append(
399                        "laterals",
400                        exp.Lateral(
401                            this=_udtf_type(unnest, has_multi_expr)(this=e),
402                            view=True,
403                            alias=exp.TableAlias(
404                                this=alias.this,  # type: ignore
405                                columns=alias_cols,
406                            ),
407                        ),
408                    )
409
410    return expression
411
412
413def explode_projection_to_unnest(
414    index_offset: int = 0,
415) -> t.Callable[[exp.Expression], exp.Expression]:
416    """Convert explode/posexplode projections into unnests."""
417
418    def _explode_projection_to_unnest(expression: exp.Expression) -> exp.Expression:
419        if isinstance(expression, exp.Select):
420            from sqlglot.optimizer.scope import Scope
421
422            taken_select_names = set(expression.named_selects)
423            taken_source_names = {name for name, _ in Scope(expression).references}
424
425            def new_name(names: t.Set[str], name: str) -> str:
426                name = find_new_name(names, name)
427                names.add(name)
428                return name
429
430            arrays: t.List[exp.Condition] = []
431            series_alias = new_name(taken_select_names, "pos")
432            series = exp.alias_(
433                exp.Unnest(
434                    expressions=[exp.GenerateSeries(start=exp.Literal.number(index_offset))]
435                ),
436                new_name(taken_source_names, "_u"),
437                table=[series_alias],
438            )
439
440            # we use list here because expression.selects is mutated inside the loop
441            for select in list(expression.selects):
442                explode = select.find(exp.Explode)
443
444                if explode:
445                    pos_alias = ""
446                    explode_alias = ""
447
448                    if isinstance(select, exp.Alias):
449                        explode_alias = select.args["alias"]
450                        alias = select
451                    elif isinstance(select, exp.Aliases):
452                        pos_alias = select.aliases[0]
453                        explode_alias = select.aliases[1]
454                        alias = select.replace(exp.alias_(select.this, "", copy=False))
455                    else:
456                        alias = select.replace(exp.alias_(select, ""))
457                        explode = alias.find(exp.Explode)
458                        assert explode
459
460                    is_posexplode = isinstance(explode, exp.Posexplode)
461                    explode_arg = explode.this
462
463                    if isinstance(explode, exp.ExplodeOuter):
464                        bracket = explode_arg[0]
465                        bracket.set("safe", True)
466                        bracket.set("offset", True)
467                        explode_arg = exp.func(
468                            "IF",
469                            exp.func(
470                                "ARRAY_SIZE", exp.func("COALESCE", explode_arg, exp.Array())
471                            ).eq(0),
472                            exp.array(bracket, copy=False),
473                            explode_arg,
474                        )
475
476                    # This ensures that we won't use [POS]EXPLODE's argument as a new selection
477                    if isinstance(explode_arg, exp.Column):
478                        taken_select_names.add(explode_arg.output_name)
479
480                    unnest_source_alias = new_name(taken_source_names, "_u")
481
482                    if not explode_alias:
483                        explode_alias = new_name(taken_select_names, "col")
484
485                        if is_posexplode:
486                            pos_alias = new_name(taken_select_names, "pos")
487
488                    if not pos_alias:
489                        pos_alias = new_name(taken_select_names, "pos")
490
491                    alias.set("alias", exp.to_identifier(explode_alias))
492
493                    series_table_alias = series.args["alias"].this
494                    column = exp.If(
495                        this=exp.column(series_alias, table=series_table_alias).eq(
496                            exp.column(pos_alias, table=unnest_source_alias)
497                        ),
498                        true=exp.column(explode_alias, table=unnest_source_alias),
499                    )
500
501                    explode.replace(column)
502
503                    if is_posexplode:
504                        expressions = expression.expressions
505                        expressions.insert(
506                            expressions.index(alias) + 1,
507                            exp.If(
508                                this=exp.column(series_alias, table=series_table_alias).eq(
509                                    exp.column(pos_alias, table=unnest_source_alias)
510                                ),
511                                true=exp.column(pos_alias, table=unnest_source_alias),
512                            ).as_(pos_alias),
513                        )
514                        expression.set("expressions", expressions)
515
516                    if not arrays:
517                        if expression.args.get("from"):
518                            expression.join(series, copy=False, join_type="CROSS")
519                        else:
520                            expression.from_(series, copy=False)
521
522                    size: exp.Condition = exp.ArraySize(this=explode_arg.copy())
523                    arrays.append(size)
524
525                    # trino doesn't support left join unnest with on conditions
526                    # if it did, this would be much simpler
527                    expression.join(
528                        exp.alias_(
529                            exp.Unnest(
530                                expressions=[explode_arg.copy()],
531                                offset=exp.to_identifier(pos_alias),
532                            ),
533                            unnest_source_alias,
534                            table=[explode_alias],
535                        ),
536                        join_type="CROSS",
537                        copy=False,
538                    )
539
540                    if index_offset != 1:
541                        size = size - 1
542
543                    expression.where(
544                        exp.column(series_alias, table=series_table_alias)
545                        .eq(exp.column(pos_alias, table=unnest_source_alias))
546                        .or_(
547                            (exp.column(series_alias, table=series_table_alias) > size).and_(
548                                exp.column(pos_alias, table=unnest_source_alias).eq(size)
549                            )
550                        ),
551                        copy=False,
552                    )
553
554            if arrays:
555                end: exp.Condition = exp.Greatest(this=arrays[0], expressions=arrays[1:])
556
557                if index_offset != 1:
558                    end = end - (1 - index_offset)
559                series.expressions[0].set("end", end)
560
561        return expression
562
563    return _explode_projection_to_unnest
564
565
566def add_within_group_for_percentiles(expression: exp.Expression) -> exp.Expression:
567    """Transforms percentiles by adding a WITHIN GROUP clause to them."""
568    if (
569        isinstance(expression, exp.PERCENTILES)
570        and not isinstance(expression.parent, exp.WithinGroup)
571        and expression.expression
572    ):
573        column = expression.this.pop()
574        expression.set("this", expression.expression.pop())
575        order = exp.Order(expressions=[exp.Ordered(this=column)])
576        expression = exp.WithinGroup(this=expression, expression=order)
577
578    return expression
579
580
581def remove_within_group_for_percentiles(expression: exp.Expression) -> exp.Expression:
582    """Transforms percentiles by getting rid of their corresponding WITHIN GROUP clause."""
583    if (
584        isinstance(expression, exp.WithinGroup)
585        and isinstance(expression.this, exp.PERCENTILES)
586        and isinstance(expression.expression, exp.Order)
587    ):
588        quantile = expression.this.this
589        input_value = t.cast(exp.Ordered, expression.find(exp.Ordered)).this
590        return expression.replace(exp.ApproxQuantile(this=input_value, quantile=quantile))
591
592    return expression
593
594
595def add_recursive_cte_column_names(expression: exp.Expression) -> exp.Expression:
596    """Uses projection output names in recursive CTE definitions to define the CTEs' columns."""
597    if isinstance(expression, exp.With) and expression.recursive:
598        next_name = name_sequence("_c_")
599
600        for cte in expression.expressions:
601            if not cte.args["alias"].columns:
602                query = cte.this
603                if isinstance(query, exp.SetOperation):
604                    query = query.this
605
606                cte.args["alias"].set(
607                    "columns",
608                    [exp.to_identifier(s.alias_or_name or next_name()) for s in query.selects],
609                )
610
611    return expression
612
613
614def epoch_cast_to_ts(expression: exp.Expression) -> exp.Expression:
615    """Replace 'epoch' in casts by the equivalent date literal."""
616    if (
617        isinstance(expression, (exp.Cast, exp.TryCast))
618        and expression.name.lower() == "epoch"
619        and expression.to.this in exp.DataType.TEMPORAL_TYPES
620    ):
621        expression.this.replace(exp.Literal.string("1970-01-01 00:00:00"))
622
623    return expression
624
625
626def eliminate_semi_and_anti_joins(expression: exp.Expression) -> exp.Expression:
627    """Convert SEMI and ANTI joins into equivalent forms that use EXIST instead."""
628    if isinstance(expression, exp.Select):
629        for join in expression.args.get("joins") or []:
630            on = join.args.get("on")
631            if on and join.kind in ("SEMI", "ANTI"):
632                subquery = exp.select("1").from_(join.this).where(on)
633                exists = exp.Exists(this=subquery)
634                if join.kind == "ANTI":
635                    exists = exists.not_(copy=False)
636
637                join.pop()
638                expression.where(exists, copy=False)
639
640    return expression
641
642
643def eliminate_full_outer_join(expression: exp.Expression) -> exp.Expression:
644    """
645    Converts a query with a FULL OUTER join to a union of identical queries that
646    use LEFT/RIGHT OUTER joins instead. This transformation currently only works
647    for queries that have a single FULL OUTER join.
648    """
649    if isinstance(expression, exp.Select):
650        full_outer_joins = [
651            (index, join)
652            for index, join in enumerate(expression.args.get("joins") or [])
653            if join.side == "FULL"
654        ]
655
656        if len(full_outer_joins) == 1:
657            expression_copy = expression.copy()
658            expression.set("limit", None)
659            index, full_outer_join = full_outer_joins[0]
660
661            tables = (expression.args["from"].alias_or_name, full_outer_join.alias_or_name)
662            join_conditions = full_outer_join.args.get("on") or exp.and_(
663                *[
664                    exp.column(col, tables[0]).eq(exp.column(col, tables[1]))
665                    for col in full_outer_join.args.get("using")
666                ]
667            )
668
669            full_outer_join.set("side", "left")
670            anti_join_clause = exp.select("1").from_(expression.args["from"]).where(join_conditions)
671            expression_copy.args["joins"][index].set("side", "right")
672            expression_copy = expression_copy.where(exp.Exists(this=anti_join_clause).not_())
673            expression_copy.args.pop("with", None)  # remove CTEs from RIGHT side
674            expression.args.pop("order", None)  # remove order by from LEFT side
675
676            return exp.union(expression, expression_copy, copy=False, distinct=False)
677
678    return expression
679
680
681def move_ctes_to_top_level(expression: E) -> E:
682    """
683    Some dialects (e.g. Hive, T-SQL, Spark prior to version 3) only allow CTEs to be
684    defined at the top-level, so for example queries like:
685
686        SELECT * FROM (WITH t(c) AS (SELECT 1) SELECT * FROM t) AS subq
687
688    are invalid in those dialects. This transformation can be used to ensure all CTEs are
689    moved to the top level so that the final SQL code is valid from a syntax standpoint.
690
691    TODO: handle name clashes whilst moving CTEs (it can get quite tricky & costly).
692    """
693    top_level_with = expression.args.get("with")
694    for inner_with in expression.find_all(exp.With):
695        if inner_with.parent is expression:
696            continue
697
698        if not top_level_with:
699            top_level_with = inner_with.pop()
700            expression.set("with", top_level_with)
701        else:
702            if inner_with.recursive:
703                top_level_with.set("recursive", True)
704
705            parent_cte = inner_with.find_ancestor(exp.CTE)
706            inner_with.pop()
707
708            if parent_cte:
709                i = top_level_with.expressions.index(parent_cte)
710                top_level_with.expressions[i:i] = inner_with.expressions
711                top_level_with.set("expressions", top_level_with.expressions)
712            else:
713                top_level_with.set(
714                    "expressions", top_level_with.expressions + inner_with.expressions
715                )
716
717    return expression
718
719
720def ensure_bools(expression: exp.Expression) -> exp.Expression:
721    """Converts numeric values used in conditions into explicit boolean expressions."""
722    from sqlglot.optimizer.canonicalize import ensure_bools
723
724    def _ensure_bool(node: exp.Expression) -> None:
725        if (
726            node.is_number
727            or (
728                not isinstance(node, exp.SubqueryPredicate)
729                and node.is_type(exp.DataType.Type.UNKNOWN, *exp.DataType.NUMERIC_TYPES)
730            )
731            or (isinstance(node, exp.Column) and not node.type)
732        ):
733            node.replace(node.neq(0))
734
735    for node in expression.walk():
736        ensure_bools(node, _ensure_bool)
737
738    return expression
739
740
741def unqualify_columns(expression: exp.Expression) -> exp.Expression:
742    for column in expression.find_all(exp.Column):
743        # We only wanna pop off the table, db, catalog args
744        for part in column.parts[:-1]:
745            part.pop()
746
747    return expression
748
749
750def remove_unique_constraints(expression: exp.Expression) -> exp.Expression:
751    assert isinstance(expression, exp.Create)
752    for constraint in expression.find_all(exp.UniqueColumnConstraint):
753        if constraint.parent:
754            constraint.parent.pop()
755
756    return expression
757
758
759def ctas_with_tmp_tables_to_create_tmp_view(
760    expression: exp.Expression,
761    tmp_storage_provider: t.Callable[[exp.Expression], exp.Expression] = lambda e: e,
762) -> exp.Expression:
763    assert isinstance(expression, exp.Create)
764    properties = expression.args.get("properties")
765    temporary = any(
766        isinstance(prop, exp.TemporaryProperty)
767        for prop in (properties.expressions if properties else [])
768    )
769
770    # CTAS with temp tables map to CREATE TEMPORARY VIEW
771    if expression.kind == "TABLE" and temporary:
772        if expression.expression:
773            return exp.Create(
774                kind="TEMPORARY VIEW",
775                this=expression.this,
776                expression=expression.expression,
777            )
778        return tmp_storage_provider(expression)
779
780    return expression
781
782
783def move_schema_columns_to_partitioned_by(expression: exp.Expression) -> exp.Expression:
784    """
785    In Hive, the PARTITIONED BY property acts as an extension of a table's schema. When the
786    PARTITIONED BY value is an array of column names, they are transformed into a schema.
787    The corresponding columns are removed from the create statement.
788    """
789    assert isinstance(expression, exp.Create)
790    has_schema = isinstance(expression.this, exp.Schema)
791    is_partitionable = expression.kind in {"TABLE", "VIEW"}
792
793    if has_schema and is_partitionable:
794        prop = expression.find(exp.PartitionedByProperty)
795        if prop and prop.this and not isinstance(prop.this, exp.Schema):
796            schema = expression.this
797            columns = {v.name.upper() for v in prop.this.expressions}
798            partitions = [col for col in schema.expressions if col.name.upper() in columns]
799            schema.set("expressions", [e for e in schema.expressions if e not in partitions])
800            prop.replace(exp.PartitionedByProperty(this=exp.Schema(expressions=partitions)))
801            expression.set("this", schema)
802
803    return expression
804
805
806def move_partitioned_by_to_schema_columns(expression: exp.Expression) -> exp.Expression:
807    """
808    Spark 3 supports both "HIVEFORMAT" and "DATASOURCE" formats for CREATE TABLE.
809
810    Currently, SQLGlot uses the DATASOURCE format for Spark 3.
811    """
812    assert isinstance(expression, exp.Create)
813    prop = expression.find(exp.PartitionedByProperty)
814    if (
815        prop
816        and prop.this
817        and isinstance(prop.this, exp.Schema)
818        and all(isinstance(e, exp.ColumnDef) and e.kind for e in prop.this.expressions)
819    ):
820        prop_this = exp.Tuple(
821            expressions=[exp.to_identifier(e.this) for e in prop.this.expressions]
822        )
823        schema = expression.this
824        for e in prop.this.expressions:
825            schema.append("expressions", e)
826        prop.set("this", prop_this)
827
828    return expression
829
830
831def struct_kv_to_alias(expression: exp.Expression) -> exp.Expression:
832    """Converts struct arguments to aliases, e.g. STRUCT(1 AS y)."""
833    if isinstance(expression, exp.Struct):
834        expression.set(
835            "expressions",
836            [
837                exp.alias_(e.expression, e.this) if isinstance(e, exp.PropertyEQ) else e
838                for e in expression.expressions
839            ],
840        )
841
842    return expression
843
844
845def eliminate_join_marks(expression: exp.Expression) -> exp.Expression:
846    """
847    Remove join marks from an AST. This rule assumes that all marked columns are qualified.
848    If this does not hold for a query, consider running `sqlglot.optimizer.qualify` first.
849
850    For example,
851        SELECT * FROM a, b WHERE a.id = b.id(+)    -- ... is converted to
852        SELECT * FROM a LEFT JOIN b ON a.id = b.id -- this
853
854    Args:
855        expression: The AST to remove join marks from.
856
857    Returns:
858       The AST with join marks removed.
859    """
860    from sqlglot.optimizer.scope import traverse_scope
861
862    for scope in traverse_scope(expression):
863        query = scope.expression
864
865        where = query.args.get("where")
866        joins = query.args.get("joins")
867
868        if not where or not joins:
869            continue
870
871        query_from = query.args["from"]
872
873        # These keep track of the joins to be replaced
874        new_joins: t.Dict[str, exp.Join] = {}
875        old_joins = {join.alias_or_name: join for join in joins}
876
877        for column in scope.columns:
878            if not column.args.get("join_mark"):
879                continue
880
881            predicate = column.find_ancestor(exp.Predicate, exp.Select)
882            assert isinstance(
883                predicate, exp.Binary
884            ), "Columns can only be marked with (+) when involved in a binary operation"
885
886            predicate_parent = predicate.parent
887            join_predicate = predicate.pop()
888
889            left_columns = [
890                c for c in join_predicate.left.find_all(exp.Column) if c.args.get("join_mark")
891            ]
892            right_columns = [
893                c for c in join_predicate.right.find_all(exp.Column) if c.args.get("join_mark")
894            ]
895
896            assert not (
897                left_columns and right_columns
898            ), "The (+) marker cannot appear in both sides of a binary predicate"
899
900            marked_column_tables = set()
901            for col in left_columns or right_columns:
902                table = col.table
903                assert table, f"Column {col} needs to be qualified with a table"
904
905                col.set("join_mark", False)
906                marked_column_tables.add(table)
907
908            assert (
909                len(marked_column_tables) == 1
910            ), "Columns of only a single table can be marked with (+) in a given binary predicate"
911
912            # Add predicate if join already copied, or add join if it is new
913            join_this = old_joins.get(col.table, query_from).this
914            existing_join = new_joins.get(join_this.alias_or_name)
915            if existing_join:
916                existing_join.set("on", exp.and_(existing_join.args["on"], join_predicate))
917            else:
918                new_joins[join_this.alias_or_name] = exp.Join(
919                    this=join_this.copy(), on=join_predicate.copy(), kind="LEFT"
920                )
921
922            # If the parent of the target predicate is a binary node, then it now has only one child
923            if isinstance(predicate_parent, exp.Binary):
924                if predicate_parent.left is None:
925                    predicate_parent.replace(predicate_parent.right)
926                else:
927                    predicate_parent.replace(predicate_parent.left)
928
929        if query_from.alias_or_name in new_joins:
930            only_old_joins = old_joins.keys() - new_joins.keys()
931            assert (
932                len(only_old_joins) >= 1
933            ), "Cannot determine which table to use in the new FROM clause"
934
935            new_from_name = list(only_old_joins)[0]
936            query.set("from", exp.From(this=old_joins[new_from_name].this))
937
938        if new_joins:
939            query.set("joins", list(new_joins.values()))
940
941        if not where.this:
942            where.pop()
943
944    return expression
945
946
947def any_to_exists(expression: exp.Expression) -> exp.Expression:
948    """
949    Transform ANY operator to Spark's EXISTS
950
951    For example,
952        - Postgres: SELECT * FROM tbl WHERE 5 > ANY(tbl.col)
953        - Spark: SELECT * FROM tbl WHERE EXISTS(tbl.col, x -> x < 5)
954
955    Both ANY and EXISTS accept queries but currently only array expressions are supported for this
956    transformation
957    """
958    if isinstance(expression, exp.Select):
959        for any in expression.find_all(exp.Any):
960            this = any.this
961            if isinstance(this, exp.Query):
962                continue
963
964            binop = any.parent
965            if isinstance(binop, exp.Binary):
966                lambda_arg = exp.to_identifier("x")
967                any.replace(lambda_arg)
968                lambda_expr = exp.Lambda(this=binop.copy(), expressions=[lambda_arg])
969                binop.replace(exp.Exists(this=this.unnest(), expression=lambda_expr))
970
971    return expression