Merge branch 'main' into main_chelsealin_perf

Author: chelsea-lin

bigframes/core/compile/sqlglot/compiler.py

@@ -153,13 +153,17 @@ def compile_sql_select(node: sql_nodes.SqlSelectNode, child: ir.SQLGlotIR):
         for ordering in node.sorting
     )
 
-    projected_cols: tuple[tuple[str, sge.Expression], ...] = tuple(
-        (
-            cdef.id.sql,
-            expression_compiler.expression_compiler.compile_expression(cdef.expression),
+    projected_cols: tuple[tuple[str, sge.Expression], ...] = tuple()
+    if not node.is_star_selection:
+        projected_cols = tuple(
+            (
+                cdef.id.sql,
+                expression_compiler.expression_compiler.compile_expression(
+                    cdef.expression
+                ),
+            )
+            for cdef in node.selections
         )
-        for cdef in node.selections
-    )
 
     sge_predicates = tuple(
         expression_compiler.expression_compiler.compile_expression(expression)

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -150,7 +150,7 @@ def from_table(
         if sql_predicate:
             select_expr = sge.Select().select(sge.Star()).from_(table_expr)
             select_expr = select_expr.where(
-                sg.parse_one(sql_predicate, dialect="bigquery"), append=False
+                sg.parse_one(sql_predicate, dialect=cls.dialect), append=False
             )
             return cls(expr=select_expr, uid_gen=uid_gen)
 
@@ -172,16 +172,19 @@ def select(
         if len(sorting) > 0:
             new_expr = new_expr.order_by(*sorting)
 
-        to_select = [
-            sge.Alias(
-                this=expr,
-                alias=sge.to_identifier(id, quoted=self.quoted),
-            )
-            if expr.alias_or_name != id
-            else expr
-            for id, expr in selections
-        ]
-        new_expr = new_expr.select(*to_select, append=False)
+        if len(selections) > 0:
+            to_select = [
+                sge.Alias(
+                    this=expr,
+                    alias=sge.to_identifier(id, quoted=self.quoted),
+                )
+                if expr.alias_or_name != id
+                else expr
+                for id, expr in selections
+            ]
+            new_expr = new_expr.select(*to_select, append=False)
+        else:
+            new_expr = new_expr.select(sge.Star(), append=False)
 
         if len(predicates) > 0:
             condition = _and(predicates)

bigframes/core/rewrite/select_pullup.py

@@ -54,13 +54,12 @@ 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
+                nodes.ScanItem(
+                    identifiers.ColumnId(scan_item.source_id), scan_item.source_id
+                )
+                for scan_item in node.scan_list.items
             ]
         )
         new_source = dataclasses.replace(node, scan_list=new_scan_list)

bigframes/core/sql_nodes.py

@@ -142,6 +142,10 @@ def consumed_ids(self):
     def _node_expressions(self):
         raise NotImplementedError()
 
+    @property
+    def is_star_selection(self) -> bool:
+        return tuple(self.ids) == tuple(self.child.ids)
+
     @functools.cache
     def get_id_mapping(self) -> dict[identifiers.ColumnId, ex.Expression]:
         return {cdef.id: cdef.expression for cdef in self.selections}

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

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

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

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

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

@@ -1,20 +1,6 @@
 WITH `bfcte_0` AS (
   SELECT
-    `bool_col`,
-    `bytes_col`,
-    `date_col`,
-    `datetime_col`,
-    `duration_col`,
-    `float64_col`,
-    `geography_col`,
-    `int64_col`,
-    `int64_too`,
-    `numeric_col`,
-    `rowindex`,
-    `rowindex_2`,
-    `string_col`,
-    `time_col`,
-    `timestamp_col`
+    *
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT

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

@@ -1,8 +1,8 @@
 WITH `bfcte_0` AS (
   SELECT
     `bool_col`,
-    `duration_col`,
     `int64_col`,
+    `duration_col`,
     `int64_col` AS `bfcol_6`,
     `bool_col` AS `bfcol_7`,
     `duration_col` AS `bfcol_8`

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

@@ -1,8 +1,8 @@
 WITH `bfcte_0` AS (
   SELECT
     `bool_col`,
-    `duration_col`,
     `int64_col`,
+    `duration_col`,
     `int64_col` AS `bfcol_6`,
     `bool_col` AS `bfcol_7`,
     `duration_col` AS `bfcol_8`

tests/unit/core/compile/sqlglot/snapshots/test_compile_explode/test_compile_explode_dataframe/out.sql

@@ -1,7 +1,7 @@
 WITH `bfcte_0` AS (
   SELECT
-    `int_list_col`,
     `rowindex`,
+    `int_list_col`,
     `string_list_col`
   FROM `bigframes-dev`.`sqlglot_test`.`repeated_types`
 ), `bfcte_1` AS (