Skip to content
This repository was archived by the owner on Apr 1, 2026. It is now read-only.

Commit dd8825c

Browse files
committed
refactor: enable SELECT * optimizations in sqlglot compiler
1 parent 3e07c80 commit dd8825c

File tree

7 files changed

+32
-52
lines changed
  • bigframes/core
  • tests/unit/core/compile/sqlglot
    • aggregations/snapshots/test_nullary_compiler/test_size
    • snapshots/test_compile_readtable
      • test_compile_readtable_w_columns_filters
      • test_compile_readtable_w_json_types
      • test_compile_readtable_w_system_time

7 files changed

+32
-52
lines changed

bigframes/core/compile/sqlglot/compiler.py

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -153,13 +153,17 @@ def compile_sql_select(node: sql_nodes.SqlSelectNode, child: ir.SQLGlotIR):
153153
for ordering in node.sorting
154154
)
155155

156-
projected_cols: tuple[tuple[str, sge.Expression], ...] = tuple(
157-
(
158-
cdef.id.sql,
159-
expression_compiler.expression_compiler.compile_expression(cdef.expression),
156+
projected_cols: tuple[tuple[str, sge.Expression], ...] = tuple()
157+
if not node.is_star_selection:
158+
projected_cols = tuple(
159+
(
160+
cdef.id.sql,
161+
expression_compiler.expression_compiler.compile_expression(
162+
cdef.expression
163+
),
164+
)
165+
for cdef in node.selections
160166
)
161-
for cdef in node.selections
162-
)
163167

164168
sge_predicates = tuple(
165169
expression_compiler.expression_compiler.compile_expression(expression)

bigframes/core/compile/sqlglot/sqlglot_ir.py

Lines changed: 14 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -150,7 +150,7 @@ def from_table(
150150
if sql_predicate:
151151
select_expr = sge.Select().select(sge.Star()).from_(table_expr)
152152
select_expr = select_expr.where(
153-
sg.parse_one(sql_predicate, dialect="bigquery"), append=False
153+
sg.parse_one(sql_predicate, dialect=cls.dialect), append=False
154154
)
155155
return cls(expr=select_expr, uid_gen=uid_gen)
156156

@@ -172,16 +172,19 @@ def select(
172172
if len(sorting) > 0:
173173
new_expr = new_expr.order_by(*sorting)
174174

175-
to_select = [
176-
sge.Alias(
177-
this=expr,
178-
alias=sge.to_identifier(id, quoted=self.quoted),
179-
)
180-
if expr.alias_or_name != id
181-
else expr
182-
for id, expr in selections
183-
]
184-
new_expr = new_expr.select(*to_select, append=False)
175+
if len(selections) > 0:
176+
to_select = [
177+
sge.Alias(
178+
this=expr,
179+
alias=sge.to_identifier(id, quoted=self.quoted),
180+
)
181+
if expr.alias_or_name != id
182+
else expr
183+
for id, expr in selections
184+
]
185+
new_expr = new_expr.select(*to_select, append=False)
186+
else:
187+
new_expr = new_expr.select(sge.Star(), append=False)
185188

186189
if len(predicates) > 0:
187190
condition = _and(predicates)

bigframes/core/sql_nodes.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,10 @@ def consumed_ids(self):
142142
def _node_expressions(self):
143143
raise NotImplementedError()
144144

145+
@property
146+
def is_star_selection(self) -> bool:
147+
return tuple(self.ids) == tuple(self.child.ids)
148+
145149
@functools.cache
146150
def get_id_mapping(self) -> dict[identifiers.ColumnId, ex.Expression]:
147151
return {cdef.id: cdef.expression for cdef in self.selections}

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

Lines changed: 1 addition & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,6 @@
11
WITH `bfcte_0` AS (
22
SELECT
3-
`bool_col`,
4-
`bytes_col`,
5-
`date_col`,
6-
`datetime_col`,
7-
`geography_col`,
8-
`int64_col`,
9-
`int64_too`,
10-
`numeric_col`,
11-
`float64_col`,
12-
`rowindex`,
13-
`rowindex_2`,
14-
`string_col`,
15-
`time_col`,
16-
`timestamp_col`,
17-
`duration_col`
3+
*
184
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
195
), `bfcte_1` AS (
206
SELECT

tests/unit/core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_readtable_w_columns_filters/out.sql

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,5 @@ WITH `bfcte_0` AS (
66
`rowindex` > 0 AND `string_col` IN ('Hello, World!')
77
)
88
SELECT
9-
`rowindex`,
10-
`int64_col`,
11-
`string_col`
9+
*
1210
FROM `bfcte_0`
Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
11
SELECT
2-
`rowindex`,
3-
`json_col`
2+
*
43
FROM `bigframes-dev`.`sqlglot_test`.`json_types`
Lines changed: 1 addition & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,3 @@
11
SELECT
2-
`bool_col`,
3-
`bytes_col`,
4-
`date_col`,
5-
`datetime_col`,
6-
`geography_col`,
7-
`int64_col`,
8-
`int64_too`,
9-
`numeric_col`,
10-
`float64_col`,
11-
`rowindex`,
12-
`rowindex_2`,
13-
`string_col`,
14-
`time_col`,
15-
`timestamp_col`,
16-
`duration_col`
2+
*
173
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` FOR SYSTEM_TIME AS OF '2025-11-09T03:04:05.678901+00:00'

0 commit comments

Comments
 (0)