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