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