refactor: add table aliases to prevent ambiguity in SQLGlotIR

Author: chelsea-lin

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -134,24 +134,50 @@ def from_table(
         """
         version = (
             sge.Version(
-                this="TIMESTAMP",
-                expression=sge.Literal(this=system_time.isoformat(), is_string=True),
+                this=sge.Identifier(this="SYSTEM_TIME", quoted=False),
+                expression=sge.Literal.string(system_time.isoformat()),
                 kind="AS OF",
             )
             if system_time
             else None
         )
+        table_alias = next(uid_gen.get_uid_stream("bft_"))
         table_expr = sge.Table(
             this=sg.to_identifier(table_id, quoted=cls.quoted),
             db=sg.to_identifier(dataset_id, quoted=cls.quoted),
             catalog=sg.to_identifier(project_id, quoted=cls.quoted),
             version=version,
+            alias=sge.Identifier(this=table_alias, quoted=cls.quoted),
         )
         if sql_predicate:
-            select_expr = sge.Select().select(sge.Star()).from_(table_expr)
-            select_expr = select_expr.where(
-                sg.parse_one(sql_predicate, dialect=cls.dialect), append=False
+            table_alias = sge.to_identifier(
+                next(uid_gen.get_uid_stream("bft_")), quoted=cls.quoted
             )
+            # WORKAROUND: SQLGlot renders Table + version + alias in wrong order for BigQuery.
+            # Wrapping in a subquery ensures valid SQL: (SELECT * FROM table FOR SYSTEM_TIME AS OF ...) AS alias
+            if version:
+                from_item = (
+                    sge.Select()
+                    .select(sge.Star())
+                    .from_(table_expr)
+                    .subquery(alias=table_alias)
+                )
+            else:
+                from_item = sge.Alias(this=table_expr, alias=table_alias)
+
+            select_expr = (
+                sge.Select()
+                .select(sge.Column(this=sge.Star(), table=table_alias))
+                .from_(from_item)
+            )
+
+            predicate_expr = sg.parse_one(sql_predicate, dialect=cls.dialect)
+            predicate_expr = predicate_expr.transform(
+                lambda e: sge.Column(this=e.this, table=table_alias)
+                if isinstance(e, sge.Column) and not e.table
+                else e
+            )
+            select_expr = select_expr.where(predicate_expr, append=False)
             return cls(expr=select_expr, uid_gen=uid_gen)
 
         return cls(expr=table_expr, uid_gen=uid_gen)
@@ -165,35 +191,66 @@ def select(
     ) -> SQLGlotIR:
         # TODO: Explicitly insert CTEs into plan
         if isinstance(self.expr, sge.Select):
-            new_expr, _ = self._select_to_cte()
+            new_expr, table_alias = self._select_to_cte()
         else:
-            new_expr = sge.Select().from_(self.expr)
+            table_alias = sge.to_identifier(
+                next(self.uid_gen.get_uid_stream("bft_")), quoted=self.quoted
+            )
+            # WORKAROUND: SQLGlot renders Table + version + alias in wrong order for BigQuery.
+            # Wrapping in a subquery ensures valid SQL: (SELECT * FROM table FOR SYSTEM_TIME AS OF ...) AS alias
+            if isinstance(self.expr, sge.Table) and self.expr.args.get("version"):
+                from_item = (
+                    sge.Select()
+                    .select(sge.Star())
+                    .from_(self.expr)
+                    .subquery(alias=table_alias)
+                )
+            else:
+                from_item = sge.Alias(this=self.expr, alias=table_alias)
+
+            new_expr = sge.Select().from_(from_item)
 
         if len(sorting) > 0:
-            new_expr = new_expr.order_by(*sorting)
+            new_expr = new_expr.order_by(
+                *[self._qualify(sort, table_alias) for sort in sorting]
+            )
 
         if len(selections) > 0:
             to_select = [
                 sge.Alias(
-                    this=expr,
+                    this=self._qualify(expr, table_alias),
                     alias=sge.to_identifier(id, quoted=self.quoted),
                 )
                 if expr.alias_or_name != id
-                else expr
+                else self._qualify(expr, table_alias)
                 for id, expr in selections
             ]
             new_expr = new_expr.select(*to_select, append=False)
         else:
-            new_expr = new_expr.select(sge.Star(), append=False)
+            new_expr = new_expr.select(
+                sge.Column(this=sge.Star(), table=table_alias), append=False
+            )
 
         if len(predicates) > 0:
-            condition = _and(predicates)
+            condition = _and(
+                tuple(self._qualify(predicate, table_alias) for predicate in predicates)
+            )
             new_expr = new_expr.where(condition, append=False)
         if limit is not None:
             new_expr = new_expr.limit(limit)
 
         return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
 
+    def _qualify(
+        self, expr: sge.Expression, table_alias: sge.Identifier
+    ) -> sge.Expression:
+        def _transform(e):
+            if isinstance(e, sge.Column) and not e.table:
+                return sge.Column(this=e.this, table=table_alias)
+            return e
+
+        return expr.transform(_transform)
+
     @classmethod
     def from_query_string(
         cls,
@@ -210,7 +267,11 @@ def from_query_string(
             this=query_string,
             alias=cte_name,
         )
-        select_expr = sge.Select().select(sge.Star()).from_(sge.Table(this=cte_name))
+        select_expr = (
+            sge.Select()
+            .select(sge.Column(this=sge.Star(), table=cte_name))
+            .from_(sge.Table(this=cte_name))
+        )
         select_expr = _set_query_ctes(select_expr, [cte])
         return cls(expr=select_expr, uid_gen=uid_gen)
 
@@ -276,9 +337,23 @@ def join(
         right_select, right_ctes = _pop_query_ctes(right_select)
         merged_ctes = _merge_ctes(left_ctes, right_ctes)
 
+        # Qualify join conditions
+        qualified_conditions = tuple(
+            (
+                typed_expr.TypedExpr(
+                    self._qualify(left.expr, left_cte_name), left.dtype
+                ),
+                typed_expr.TypedExpr(
+                    self._qualify(right.expr, right_cte_name), right.dtype
+                ),
+            )
+            for left, right in conditions
+        )
+
         join_on = _and(
             tuple(
-                _join_condition(left, right, joins_nulls) for left, right in conditions
+                _join_condition(left, right, joins_nulls)
+                for left, right in qualified_conditions
             )
         )
 
@@ -310,7 +385,7 @@ def isin_join(
         merged_ctes = _merge_ctes(left_ctes, right_ctes)
 
         left_condition = typed_expr.TypedExpr(
-            sge.Column(this=conditions[0].expr, table=left_cte_name),
+            self._qualify(conditions[0].expr, left_cte_name),
             conditions[0].dtype,
         )
 
@@ -320,7 +395,7 @@ def isin_join(
                 next(self.uid_gen.get_uid_stream("bft_")), quoted=self.quoted
             )
             right_condition = typed_expr.TypedExpr(
-                sge.Column(this=conditions[1].expr, table=right_table_name),
+                self._qualify(conditions[1].expr, right_table_name),
                 conditions[1].dtype,
             )
             new_column = sge.Exists(
@@ -371,7 +446,8 @@ def sample(self, fraction: float) -> SQLGlotIR:
             expression=_literal(fraction, dtypes.FLOAT_DTYPE),
         )
 
-        new_expr = self._select_to_cte()[0].where(condition, append=False)
+        new_expr, table_alias = self._select_to_cte()
+        new_expr = new_expr.where(condition, append=False)
         return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
 
     def aggregate(
@@ -387,23 +463,27 @@ def aggregate(
             by_cols: column expressions for aggregation
             dropna_cols: columns whether null keys should be dropped
         """
-        aggregations_expr = [
+        new_expr, table_alias = self._select_to_cte()
+
+        qualified_aggregations_expr = [
             sge.Alias(
-                this=expr,
+                this=self._qualify(expr, table_alias),
                 alias=sge.to_identifier(id, quoted=self.quoted),
             )
             for id, expr in aggregations
         ]
 
-        new_expr, _ = self._select_to_cte()
-        new_expr = new_expr.group_by(*by_cols).select(
-            *[*by_cols, *aggregations_expr], append=False
+        qualified_by_cols = [self._qualify(col, table_alias) for col in by_cols]
+        qualified_dropna_cols = [self._qualify(col, table_alias) for col in dropna_cols]
+
+        new_expr = new_expr.group_by(*qualified_by_cols).select(
+            *[*qualified_by_cols, *qualified_aggregations_expr], append=False
         )
 
         condition = _and(
             tuple(
                 sg.not_(sge.Is(this=drop_col, expression=sge.Null()))
-                for drop_col in dropna_cols
+                for drop_col in qualified_dropna_cols
             )
         )
         if condition is not None:
@@ -496,14 +576,16 @@ def _explode_single_column(
         unnested_column_alias = sge.to_identifier(
             next(self.uid_gen.get_uid_stream("bfcol_")), quoted=self.quoted
         )
+
+        new_expr, table_alias = self._select_to_cte()
+
         unnest_expr = sge.Unnest(
-            expressions=[column],
+            expressions=[sge.Column(this=column, table=table_alias)],
             alias=sge.TableAlias(columns=[unnested_column_alias]),
             offset=offset,
         )
         selection = sge.Star(replace=[unnested_column_alias.as_(column)])
 
-        new_expr, _ = self._select_to_cte()
         # Use LEFT JOIN to preserve rows when unnesting empty arrays.
         new_expr = new_expr.select(selection, append=False).join(
             unnest_expr, join_type="LEFT"
@@ -524,10 +606,12 @@ def _explode_multiple_columns(
             for column_name in column_names
         ]
 
+        new_expr, table_alias = self._select_to_cte()
+
         # If there are multiple columns, we need to unnest by zipping the arrays:
         # https://cloud.google.com/bigquery/docs/arrays#zipping_arrays
         column_lengths = [
-            sge.func("ARRAY_LENGTH", sge.to_identifier(column, quoted=self.quoted)) - 1
+            sge.func("ARRAY_LENGTH", sge.Column(this=column, table=table_alias)) - 1
             for column in columns
         ]
         generate_array = sge.func(
@@ -554,7 +638,6 @@ def _explode_multiple_columns(
                 for column in columns
             ]
         )
-        new_expr, _ = self._select_to_cte()
         # Use LEFT JOIN to preserve rows when unnesting empty arrays.
         new_expr = new_expr.select(selection, append=False).join(
             unnest_expr, join_type="LEFT"
@@ -590,7 +673,9 @@ def _select_to_cte(self) -> tuple[sge.Select, sge.Identifier]:
             alias=cte_name,
         )
         new_select_expr = (
-            sge.Select().select(sge.Star()).from_(sge.Table(this=cte_name))
+            sge.Select()
+            .select(sge.Column(this=sge.Star(), table=cte_name))
+            .from_(sge.Table(this=cte_name))
         )
         new_select_expr = _set_query_ctes(new_select_expr, [*existing_ctes, new_cte])
         return new_select_expr, cte_name

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_binary_compiler/test_corr/out.sql

@@ -2,7 +2,7 @@ WITH `bfcte_0` AS (
   SELECT
     `int64_col`,
     `float64_col`
-  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
 ), `bfcte_1` AS (
   SELECT
     CORR(`int64_col`, `float64_col`) AS `bfcol_2`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_binary_compiler/test_cov/out.sql

@@ -2,7 +2,7 @@ WITH `bfcte_0` AS (
   SELECT
     `int64_col`,
     `float64_col`
-  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
 ), `bfcte_1` AS (
   SELECT
     COVAR_SAMP(`int64_col`, `float64_col`) AS `bfcol_2`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_nullary_compiler/test_row_number/out.sql

@@ -1,3 +1,3 @@
 SELECT
   ROW_NUMBER() OVER () - 1 AS `row_number`
-FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_nullary_compiler/test_row_number_with_window/out.sql

@@ -1,3 +1,3 @@
 SELECT
   ROW_NUMBER() OVER (ORDER BY `int64_col` ASC NULLS LAST) - 1 AS `row_number`
-FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_nullary_compiler/test_size/out.sql

@@ -1,7 +1,7 @@
 WITH `bfcte_0` AS (
   SELECT
     *
-  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
 ), `bfcte_1` AS (
   SELECT
     COUNT(1) AS `bfcol_32`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_ordered_unary_compiler/test_array_agg/out.sql

@@ -1,7 +1,7 @@
 WITH `bfcte_0` AS (
   SELECT
     `int64_col`
-  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
 ), `bfcte_1` AS (
   SELECT
     ARRAY_AGG(`int64_col` IGNORE NULLS ORDER BY `int64_col` IS NULL ASC, `int64_col` ASC) AS `bfcol_1`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_ordered_unary_compiler/test_string_agg/out.sql

@@ -1,7 +1,7 @@
 WITH `bfcte_0` AS (
   SELECT
     `string_col`
-  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
 ), `bfcte_1` AS (
   SELECT
     COALESCE(

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_all/out.sql

@@ -2,7 +2,7 @@ WITH `bfcte_0` AS (
   SELECT
     `bool_col`,
     `int64_col`
-  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+  FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
 ), `bfcte_1` AS (
   SELECT
     COALESCE(LOGICAL_AND(`bool_col`), TRUE) AS `bfcol_2`,

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_all_w_window/out.sql

@@ -1,3 +1,3 @@
 SELECT
   COALESCE(LOGICAL_AND(`bool_col`) OVER (), TRUE) AS `agg_bool`
-FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
+FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`