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