Skip to content

fix: preserve aliases on cast columns and fix star selection in sqlglot (#17394) #17455

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
shuoweil merged 5 commits into from
Jun 12, 2026
Merged

fix: preserve aliases on cast columns and fix star selection in sqlglot (#17394) #17455

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
899d5e8
fix: preserve aliases on cast columns in sqlglot
shuoweil Jun 10, 2026
e03f408
fix: snapshot newline mismatch
shuoweil Jun 12, 2026
3d92d4d
style: format test_compile_readtable
shuoweil Jun 12, 2026
5cb3f13
fix: avoid nested aliases in sqlglot
shuoweil Jun 12, 2026
c04db27
fix: revert pre-commit config changes and append newlines to test sql…
shuoweil Jun 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions packages/bigframes/bigframes/core/compile/sqlglot/sqlglot_ir.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -249,12 +249,13 @@ def select(
# TODO: Explicitly insert CTEs into plan
if len(selections) > 0:
to_select = [
sge.Alias(
this=expr,
expr
if (isinstance(expr, sge.Alias) and expr.alias == id)
or (isinstance(expr, sge.Column) and expr.name == id)
else sge.Alias(
this=expr.this if isinstance(expr, sge.Alias) else expr,
alias=sql.identifier(id),
)
if expr.alias_or_name != id
else expr
for id, expr in selections
]
new_expr = self.expr.select(*to_select)
Expand Down
9 changes: 8 additions & 1 deletion packages/bigframes/bigframes/core/sql_nodes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -276,7 +276,14 @@ def _node_expressions(self):

@property
def is_star_selection(self) -> bool:
return tuple(self.ids) == tuple(self.child.ids)
if tuple(self.ids) != tuple(self.child.ids):
return False
for cdef in self.selections:
if not isinstance(cdef.expression, ex.DerefOp):
return False
if cdef.expression.id != cdef.id:
return False
return True

@functools.cache
def get_id_mapping(self) -> dict[identifiers.ColumnId, ex.Expression]:
Expand Down
6 changes: 3 additions & 3 deletions ...nit/core/compile/sqlglot/expressions/snapshots/test_datetime_ops/test_to_datetime/out.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
SELECT
CAST(TIMESTAMP_MICROS(CAST(TRUNC(`int64_col` * 0.001) AS INT64)) AS DATETIME) AS `int64_col`,
SAFE_CAST(`string_col` AS DATETIME),
SAFE_CAST(`string_col` AS DATETIME) AS `string_col`,
CAST(TIMESTAMP_MICROS(CAST(TRUNC(`float64_col` * 0.001) AS INT64)) AS DATETIME) AS `float64_col`,
SAFE_CAST(`timestamp_col` AS DATETIME),
SAFE_CAST(`timestamp_col` AS DATETIME) AS `timestamp_col`,
SAFE_CAST(`string_col` AS DATETIME) AS `string_col_fmt`
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
4 changes: 2 additions & 2 deletions ...nit/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_astype_float/out.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
SELECT
CAST(CAST(`bool_col` AS INT64) AS FLOAT64),
CAST(CAST(`bool_col` AS INT64) AS FLOAT64) AS `bool_col`,
CAST('1.34235e4' AS FLOAT64) AS `str_const`,
SAFE_CAST(SAFE_CAST(`bool_col` AS INT64) AS FLOAT64) AS `bool_w_safe`
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
4 changes: 2 additions & 2 deletions ...it/core/compile/sqlglot/expressions/snapshots/test_generic_ops/test_astype_string/out.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
SELECT
CAST(`int64_col` AS STRING),
CAST(`int64_col` AS STRING) AS `int64_col`,
INITCAP(CAST(`bool_col` AS STRING)) AS `bool_col`,
INITCAP(SAFE_CAST(`bool_col` AS STRING)) AS `bool_w_safe`
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
2 changes: 1 addition & 1 deletion packages/bigframes/tests/unit/core/compile/sqlglot/expressions/test_datetime_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -217,7 +217,7 @@ def test_to_datetime(scalar_types_df: bpd.DataFrame, snapshot):
)

sql = utils._apply_ops_to_sql(bf_df, list(ops_map.values()), list(ops_map.keys()))
snapshot.assert_match(sql, "out.sql")
snapshot.assert_match(sql + "\n", "out.sql")


def test_to_timestamp(scalar_types_df: bpd.DataFrame, snapshot):
Expand Down
4 changes: 2 additions & 2 deletions packages/bigframes/tests/unit/core/compile/sqlglot/expressions/test_generic_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ def test_astype_float(scalar_types_df: bpd.DataFrame, snapshot):
"bool_w_safe": ops.AsTypeOp(to_type=to_type, safe=True).as_expr("bool_col"),
}
sql = utils._apply_ops_to_sql(bf_df, list(ops_map.values()), list(ops_map.keys()))
snapshot.assert_match(sql, "out.sql")
snapshot.assert_match(sql + "\n", "out.sql")


def test_astype_bool(scalar_types_df: bpd.DataFrame, snapshot):
Expand Down Expand Up @@ -107,7 +107,7 @@ def test_astype_string(scalar_types_df: bpd.DataFrame, snapshot):
"bool_w_safe": ops.AsTypeOp(to_type=to_type, safe=True).as_expr("bool_col"),
}
sql = utils._apply_ops_to_sql(bf_df, list(ops_map.values()), list(ops_map.keys()))
snapshot.assert_match(sql, "out.sql")
snapshot.assert_match(sql + "\n", "out.sql")


def test_astype_json(scalar_types_df: bpd.DataFrame, snapshot):
Expand Down
5 changes: 5 additions & 0 deletions ...core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_astype_aliases/out.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
SELECT
`rowindex`,
CAST(`timestamp_col` AS STRING) AS `timestamp_col`,
CAST(`int64_col` AS FLOAT64) AS `int64_col`
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
12 changes: 12 additions & 0 deletions packages/bigframes/tests/unit/core/compile/sqlglot/test_compile_readtable.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,3 +80,15 @@ def test_compile_readtable_w_columns_filters(compiler_session, snapshot):
filters=filters,
)
snapshot.assert_match(bf_df.sql, "out.sql")


def test_compile_astype_aliases(scalar_types_df: bpd.DataFrame, snapshot):
# Test case for issue #17394 (CAST columns lose their aliases)
bf_df = scalar_types_df[["timestamp_col", "int64_col"]]
result = bf_df.astype(
{
"timestamp_col": "string[pyarrow]",
"int64_col": "Float64",
}
)
snapshot.assert_match(result.sql + "\n", "out.sql")
Loading