Merge branch 'main' into shuowei-anywidget-index-testcase

Author: shuoweil

bigframes/core/compile/sqlglot/__init__.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 from __future__ import annotations
 
-from bigframes.core.compile.sqlglot.compiler import SQLGlotCompiler
+from bigframes.core.compile.sqlglot.compiler import compile_sql
 import bigframes.core.compile.sqlglot.expressions.ai_ops  # noqa: F401
 import bigframes.core.compile.sqlglot.expressions.array_ops  # noqa: F401
 import bigframes.core.compile.sqlglot.expressions.blob_ops  # noqa: F401
@@ -29,4 +29,4 @@
 import bigframes.core.compile.sqlglot.expressions.struct_ops  # noqa: F401
 import bigframes.core.compile.sqlglot.expressions.timedelta_ops  # noqa: F401
 
-__all__ = ["SQLGlotCompiler"]
+__all__ = ["compile_sql"]

bigframes/core/compile/sqlglot/compiler.py

@@ -134,6 +134,8 @@ def from_table(
                 this=sge.to_identifier(col_name, quoted=cls.quoted),
                 alias=sge.to_identifier(alias_name, quoted=cls.quoted),
             )
+            if col_name != alias_name
+            else sge.to_identifier(col_name, quoted=cls.quoted)
             for col_name, alias_name in zip(col_names, alias_names)
         ]
         table_expr = sge.Table(
@@ -227,6 +229,8 @@ def select(
                 this=expr,
                 alias=sge.to_identifier(id, quoted=self.quoted),
             )
+            if expr.alias_or_name != id
+            else expr
             for id, expr in selected_cols
         ]
 

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -13,9 +13,10 @@
 # limitations under the License.
 
 import dataclasses
+import functools
 from typing import cast
 
-from bigframes.core import expression, nodes
+from bigframes.core import expression, identifiers, nodes
 
 
 def defer_selection(
@@ -26,12 +27,19 @@ def defer_selection(
 
     In many cases, these nodes will be merged or eliminated entirely, simplifying the overall tree.
     """
-    return nodes.bottom_up(root, pull_up_select)
+    return nodes.bottom_up(
+        root, functools.partial(pull_up_select, prefer_source_names=True)
+    )
 
 
-def pull_up_select(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
+def pull_up_select(
+    node: nodes.BigFrameNode, prefer_source_names: bool
+) -> nodes.BigFrameNode:
     if isinstance(node, nodes.LeafNode):
-        return node
+        if prefer_source_names and isinstance(node, nodes.ReadTableNode):
+            return pull_up_source_ids(node)
+        else:
+            return node
     if isinstance(node, nodes.JoinNode):
         return pull_up_selects_under_join(node)
     if isinstance(node, nodes.ConcatNode):
@@ -42,6 +50,32 @@ def pull_up_select(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
     return node
 
 
+def pull_up_source_ids(node: nodes.ReadTableNode) -> nodes.BigFrameNode:
+    if all(id.sql == source_id for id, source_id in node.scan_list.items):
+        return node
+    else:
+        source_ids = sorted(
+            set(scan_item.source_id for scan_item in node.scan_list.items)
+        )
+        new_scan_list = nodes.ScanList.from_items(
+            [
+                nodes.ScanItem(identifiers.ColumnId(source_id), source_id)
+                for source_id in source_ids
+            ]
+        )
+        new_source = dataclasses.replace(node, scan_list=new_scan_list)
+        new_selection = nodes.SelectionNode(
+            new_source,
+            tuple(
+                nodes.AliasedRef(
+                    expression.DerefOp(identifiers.ColumnId(source_id)), id
+                )
+                for id, source_id in node.scan_list.items
+            ),
+        )
+        return new_selection
+
+
 def pull_up_select_unary(node: nodes.UnaryNode) -> nodes.BigFrameNode:
     child = node.child
     if not isinstance(child, nodes.SelectionNode):

bigframes/core/rewrite/select_pullup.py

@@ -402,6 +402,7 @@ def get_index_cols(
     | bigframes.enums.DefaultIndexKind,
     *,
     rename_to_schema: Optional[Dict[str, str]] = None,
+    default_index_type: bigframes.enums.DefaultIndexKind = bigframes.enums.DefaultIndexKind.SEQUENTIAL_INT64,
 ) -> List[str]:
     """
     If we can get a total ordering from the table, such as via primary key
@@ -471,7 +472,11 @@ def get_index_cols(
         # find index_cols to use. This is to avoid unexpected performance and
         # resource utilization because of the default sequential index. See
         # internal issue 335727141.
-        if _is_table_clustered_or_partitioned(table) and not primary_keys:
+        if (
+            _is_table_clustered_or_partitioned(table)
+            and not primary_keys
+            and default_index_type == bigframes.enums.DefaultIndexKind.SEQUENTIAL_INT64
+        ):
             msg = bfe.format_message(
                 f"Table '{str(table.reference)}' is clustered and/or "
                 "partitioned, but BigQuery DataFrames was not able to find a "

bigframes/session/_io/bigquery/read_gbq_table.py

@@ -40,9 +40,7 @@ def __init__(
     ):
         self.bqclient = bqclient
         self._compile_fn = (
-            compile.compile_sql
-            if compiler == "ibis"
-            else sqlglot.SQLGlotCompiler()._compile_sql
+            compile.compile_sql if compiler == "ibis" else sqlglot.compile_sql
         )
         self._publisher = publisher
 

bigframes/session/direct_gbq_execution.py

@@ -696,6 +696,7 @@ def read_gbq_table(
             table=table,
             index_col=index_col,
             rename_to_schema=rename_to_schema,
+            default_index_type=self._default_index_type,
         )
         _check_index_col_param(
             index_cols,

bigframes/session/loader.py

@@ -16,15 +16,15 @@
 import typing
 
 import bigframes.core
-import bigframes.core.compile.sqlglot as sqlglot
+import bigframes.core.compile as compile
 import bigframes.session.executor
 
 
 @dataclasses.dataclass
 class SQLCompilerExecutor(bigframes.session.executor.Executor):
     """Executor for SQL compilation using sqlglot."""
 
-    compiler = sqlglot
+    compiler = compile.sqlglot
 
     def to_sql(
         self,
@@ -38,9 +38,9 @@ def to_sql(
 
         # Compared with BigQueryCachingExecutor, SQLCompilerExecutor skips
         # caching the subtree.
-        return self.compiler.SQLGlotCompiler().compile(
-            array_value.node, ordered=ordered
-        )
+        return self.compiler.compile_sql(
+            compile.CompileRequest(array_value.node, sort_rows=ordered)
+        ).sql
 
     def execute(
         self,

bigframes/testing/compiler_session.py

@@ -1,11 +1,11 @@
 WITH `bfcte_0` AS (
   SELECT
-    `int64_col` AS `bfcol_0`,
-    `float64_col` AS `bfcol_1`
+    `float64_col`,
+    `int64_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
-    CORR(`bfcol_0`, `bfcol_1`) AS `bfcol_2`
+    CORR(`int64_col`, `float64_col`) AS `bfcol_2`
   FROM `bfcte_0`
 )
 SELECT

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

@@ -1,11 +1,11 @@
 WITH `bfcte_0` AS (
   SELECT
-    `int64_col` AS `bfcol_0`,
-    `float64_col` AS `bfcol_1`
+    `float64_col`,
+    `int64_col`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
-    COVAR_SAMP(`bfcol_0`, `bfcol_1`) AS `bfcol_2`
+    COVAR_SAMP(`int64_col`, `float64_col`) AS `bfcol_2`
   FROM `bfcte_0`
 )
 SELECT