Skip to content

Commit 64995d6

Browse files
TrevorBergeronTrevorBergeron
authored
feat: Preserve source names better for more readable sql (#2243)
1 parent 2e77861 commit 64995d6

246 files changed

Lines changed: 1174 additions & 1077 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎bigframes/core/compile/sqlglot/sqlglot_ir.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,8 @@ def from_table(
134134
this=sge.to_identifier(col_name, quoted=cls.quoted),
135135
alias=sge.to_identifier(alias_name, quoted=cls.quoted),
136136
)
137+
ifcol_name!=alias_name
138+
elsesge.to_identifier(col_name, quoted=cls.quoted)
137139
forcol_name, alias_nameinzip(col_names, alias_names)
138140
]
139141
table_expr=sge.Table(
@@ -227,6 +229,8 @@ def select(
227229
this=expr,
228230
alias=sge.to_identifier(id, quoted=self.quoted),
229231
)
232+
ifexpr.alias_or_name!=id
233+
elseexpr
230234
forid, exprinselected_cols
231235
]
232236

‎bigframes/core/rewrite/select_pullup.py‎

Lines changed: 38 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,10 @@
1313
# limitations under the License.
1414

1515
importdataclasses
16+
importfunctools
1617
fromtypingimportcast
1718

18-
frombigframes.coreimportexpression, nodes
19+
frombigframes.coreimportexpression, identifiers, nodes
1920

2021

2122
defdefer_selection(
@@ -26,12 +27,19 @@ def defer_selection(
2627
2728
In many cases, these nodes will be merged or eliminated entirely, simplifying the overall tree.
2829
"""
29-
returnnodes.bottom_up(root, pull_up_select)
30+
returnnodes.bottom_up(
31+
root, functools.partial(pull_up_select, prefer_source_names=True)
32+
)
3033

3134

32-
defpull_up_select(node: nodes.BigFrameNode) ->nodes.BigFrameNode:
35+
defpull_up_select(
36+
node: nodes.BigFrameNode, prefer_source_names: bool
37+
) ->nodes.BigFrameNode:
3338
ifisinstance(node, nodes.LeafNode):
34-
returnnode
39+
ifprefer_source_namesandisinstance(node, nodes.ReadTableNode):
40+
returnpull_up_source_ids(node)
41+
else:
42+
returnnode
3543
ifisinstance(node, nodes.JoinNode):
3644
returnpull_up_selects_under_join(node)
3745
ifisinstance(node, nodes.ConcatNode):
@@ -42,6 +50,32 @@ def pull_up_select(node: nodes.BigFrameNode) -> nodes.BigFrameNode:
4250
returnnode
4351

4452

53+
defpull_up_source_ids(node: nodes.ReadTableNode) ->nodes.BigFrameNode:
54+
ifall(id.sql==source_idforid, source_idinnode.scan_list.items):
55+
returnnode
56+
else:
57+
source_ids=sorted(
58+
set(scan_item.source_idforscan_iteminnode.scan_list.items)
59+
)
60+
new_scan_list=nodes.ScanList.from_items(
61+
[
62+
nodes.ScanItem(identifiers.ColumnId(source_id), source_id)
63+
forsource_idinsource_ids
64+
]
65+
)
66+
new_source=dataclasses.replace(node, scan_list=new_scan_list)
67+
new_selection=nodes.SelectionNode(
68+
new_source,
69+
tuple(
70+
nodes.AliasedRef(
71+
expression.DerefOp(identifiers.ColumnId(source_id)), id
72+
)
73+
forid, source_idinnode.scan_list.items
74+
),
75+
)
76+
returnnew_selection
77+
78+
4579
defpull_up_select_unary(node: nodes.UnaryNode) ->nodes.BigFrameNode:
4680
child=node.child
4781
ifnotisinstance(child, nodes.SelectionNode):

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
WITH `bfcte_0`AS (
22
SELECT
3-
`int64_col`AS`bfcol_0`,
4-
`float64_col`AS`bfcol_1`
3+
`float64_col`,
4+
`int64_col`
55
FROM`bigframes-dev`.`sqlglot_test`.`scalar_types`
66
), `bfcte_1`AS (
77
SELECT
8-
CORR(`bfcol_0`, `bfcol_1`) AS`bfcol_2`
8+
CORR(`int64_col`, `float64_col`) AS`bfcol_2`
99
FROM`bfcte_0`
1010
)
1111
SELECT

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
WITH `bfcte_0`AS (
22
SELECT
3-
`int64_col`AS`bfcol_0`,
4-
`float64_col`AS`bfcol_1`
3+
`float64_col`,
4+
`int64_col`
55
FROM`bigframes-dev`.`sqlglot_test`.`scalar_types`
66
), `bfcte_1`AS (
77
SELECT
8-
COVAR_SAMP(`bfcol_0`, `bfcol_1`) AS`bfcol_2`
8+
COVAR_SAMP(`int64_col`, `float64_col`) AS`bfcol_2`
99
FROM`bfcte_0`
1010
)
1111
SELECT

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

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,20 @@
11
WITH `bfcte_0`AS (
22
SELECT
3-
`bool_col`AS`bfcol_0`,
4-
`bytes_col`AS`bfcol_1`,
5-
`date_col`AS`bfcol_2`,
6-
`datetime_col`AS`bfcol_3`,
7-
`geography_col`AS`bfcol_4`,
8-
`int64_col`AS`bfcol_5`,
9-
`int64_too`AS`bfcol_6`,
10-
`numeric_col`AS`bfcol_7`,
11-
`float64_col`AS`bfcol_8`,
12-
`rowindex`AS`bfcol_9`,
13-
`rowindex_2`AS`bfcol_10`,
14-
`string_col`AS`bfcol_11`,
15-
`time_col`AS`bfcol_12`,
16-
`timestamp_col`AS`bfcol_13`,
17-
`duration_col`AS`bfcol_14`
3+
`bool_col`,
4+
`bytes_col`,
5+
`date_col`,
6+
`datetime_col`,
7+
`duration_col`,
8+
`float64_col`,
9+
`geography_col`,
10+
`int64_col`,
11+
`int64_too`,
12+
`numeric_col`,
13+
`rowindex`,
14+
`rowindex_2`,
15+
`string_col`,
16+
`time_col`,
17+
`timestamp_col`
1818
FROM`bigframes-dev`.`sqlglot_test`.`scalar_types`
1919
), `bfcte_1`AS (
2020
SELECT

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
WITH `bfcte_0`AS (
22
SELECT
3-
`int64_col`AS`bfcol_0`
3+
`int64_col`
44
FROM`bigframes-dev`.`sqlglot_test`.`scalar_types`
55
), `bfcte_1`AS (
66
SELECT
77
*,
8-
ROW_NUMBER() OVER (ORDER BY`bfcol_0`ASC NULLS LAST) AS`bfcol_1`
8+
ROW_NUMBER() OVER (ORDER BY`int64_col`ASC NULLS LAST) AS`bfcol_1`
99
FROM`bfcte_0`
1010
)
1111
SELECT

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

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,20 @@
11
WITH `bfcte_0`AS (
22
SELECT
3-
`bool_col`AS`bfcol_0`,
4-
`bytes_col`AS`bfcol_1`,
5-
`date_col`AS`bfcol_2`,
6-
`datetime_col`AS`bfcol_3`,
7-
`geography_col`AS`bfcol_4`,
8-
`int64_col`AS`bfcol_5`,
9-
`int64_too`AS`bfcol_6`,
10-
`numeric_col`AS`bfcol_7`,
11-
`float64_col`AS`bfcol_8`,
12-
`rowindex`AS`bfcol_9`,
13-
`rowindex_2`AS`bfcol_10`,
14-
`string_col`AS`bfcol_11`,
15-
`time_col`AS`bfcol_12`,
16-
`timestamp_col`AS`bfcol_13`,
17-
`duration_col`AS`bfcol_14`
3+
`bool_col`,
4+
`bytes_col`,
5+
`date_col`,
6+
`datetime_col`,
7+
`duration_col`,
8+
`float64_col`,
9+
`geography_col`,
10+
`int64_col`,
11+
`int64_too`,
12+
`numeric_col`,
13+
`rowindex`,
14+
`rowindex_2`,
15+
`string_col`,
16+
`time_col`,
17+
`timestamp_col`
1818
FROM`bigframes-dev`.`sqlglot_test`.`scalar_types`
1919
), `bfcte_1`AS (
2020
SELECT

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
WITH `bfcte_0`AS (
22
SELECT
3-
`int64_col`AS`bfcol_0`
3+
`int64_col`
44
FROM`bigframes-dev`.`sqlglot_test`.`scalar_types`
55
), `bfcte_1`AS (
66
SELECT
7-
ARRAY_AGG(`bfcol_0` IGNORE NULLS ORDER BY`bfcol_0` IS NULLASC, `bfcol_0`ASC) AS`bfcol_1`
7+
ARRAY_AGG(`int64_col` IGNORE NULLS ORDER BY`int64_col` IS NULLASC, `int64_col`ASC) AS`bfcol_1`
88
FROM`bfcte_0`
99
)
1010
SELECT

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

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,16 @@
11
WITH `bfcte_0`AS (
22
SELECT
3-
`string_col`AS`bfcol_0`
3+
`string_col`
44
FROM`bigframes-dev`.`sqlglot_test`.`scalar_types`
55
), `bfcte_1`AS (
66
SELECT
7-
COALESCE(STRING_AGG(`bfcol_0`, ','
8-
ORDER BY
9-
`bfcol_0` IS NULLASC,
10-
`bfcol_0`ASC), '') AS`bfcol_1`
7+
COALESCE(
8+
STRING_AGG(`string_col`, ','
9+
ORDER BY
10+
`string_col` IS NULLASC,
11+
`string_col`ASC),
12+
''
13+
) AS`bfcol_1`
1114
FROM`bfcte_0`
1215
)
1316
SELECT

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
WITH `bfcte_0`AS (
22
SELECT
3-
`bool_col`AS`bfcol_0`
3+
`bool_col`
44
FROM`bigframes-dev`.`sqlglot_test`.`scalar_types`
55
), `bfcte_1`AS (
66
SELECT
7-
COALESCE(LOGICAL_AND(`bfcol_0`), TRUE) AS`bfcol_1`
7+
COALESCE(LOGICAL_AND(`bool_col`), TRUE) AS`bfcol_1`
88
FROM`bfcte_0`
99
)
1010
SELECT

0 commit comments

Comments
 (0)