Skip to content

Commit 5c2eadc

Browse files
jcristjcrist
committed
fix(sql): properly parenthesize binary ops containing named expressions
1 parent 48be246 commit 5c2eadc

File tree

15 files changed

+104
-45
lines changed

15 files changed

+104
-45
lines changed

‎ibis/backends/sql/compilers/base.py

-3
Original file line numberDiff line numberDiff line change
@@ -689,9 +689,6 @@ def visit_Cast(self, op, *, arg, to):
689689
def visit_ScalarSubquery(self, op, *, rel):
690690
return rel.this.subquery(copy=False)
691691

692-
def visit_Alias(self, op, *, arg, name):
693-
return arg
694-
695692
def visit_Literal(self, op, *, value, dtype):
696693
"""Compile a literal value.
697694

‎ibis/backends/sql/compilers/bigquery/__init__.py

+10-4
Original file line numberDiff line numberDiff line change
@@ -1017,7 +1017,14 @@ def visit_DropColumns(self, op, *, parent, columns_to_drop):
10171017
return sg.select(column).from_(parent)
10181018

10191019
def visit_TableUnnest(
1020-
self, op, *, parent, column, offset: str | None, keep_empty: bool
1020+
self,
1021+
op,
1022+
*,
1023+
parent,
1024+
column,
1025+
column_name: str,
1026+
offset: str | None,
1027+
keep_empty: bool,
10211028
):
10221029
quoted = self.quoted
10231030

@@ -1029,9 +1036,8 @@ def visit_TableUnnest(
10291036

10301037
table = sg.to_identifier(parent.alias_or_name, quoted=quoted)
10311038

1032-
opname = op.column.name
1033-
overlaps_with_parent = opname in op.parent.schema
1034-
computed_column = column_alias.as_(opname, quoted=quoted)
1039+
overlaps_with_parent = column_name in op.parent.schema
1040+
computed_column = column_alias.as_(column_name, quoted=quoted)
10351041

10361042
# replace the existing column if the unnested column hasn't been
10371043
# renamed

‎ibis/backends/sql/compilers/clickhouse.py

+10-4
Original file line numberDiff line numberDiff line change
@@ -688,7 +688,14 @@ def visit_DropColumns(self, op, *, parent, columns_to_drop):
688688
return sg.select(column).from_(parent)
689689

690690
def visit_TableUnnest(
691-
self, op, *, parent, column, offset: str | None, keep_empty: bool
691+
self,
692+
op,
693+
*,
694+
parent,
695+
column,
696+
column_name: str,
697+
offset: str | None,
698+
keep_empty: bool,
692699
):
693700
quoted = self.quoted
694701

@@ -700,9 +707,8 @@ def visit_TableUnnest(
700707

701708
selcols = []
702709

703-
opname = op.column.name
704-
overlaps_with_parent = opname in op.parent.schema
705-
computed_column = column_alias.as_(opname, quoted=quoted)
710+
overlaps_with_parent = column_name in op.parent.schema
711+
computed_column = column_alias.as_(column_name, quoted=quoted)
706712

707713
if offset is not None:
708714
if overlaps_with_parent:

‎ibis/backends/sql/compilers/duckdb.py

+13-5
Original file line numberDiff line numberDiff line change
@@ -609,15 +609,21 @@ def visit_DropColumns(self, op, *, parent, columns_to_drop):
609609
return sg.select(column).from_(parent)
610610

611611
def visit_TableUnnest(
612-
self, op, *, parent, column, offset: str | None, keep_empty: bool
612+
self,
613+
op,
614+
*,
615+
parent,
616+
column,
617+
column_name: str,
618+
offset: str | None,
619+
keep_empty: bool,
613620
):
614621
quoted = self.quoted
615622

616623
column_alias = sg.to_identifier(gen_name("table_unnest_column"), quoted=quoted)
617624

618-
opname = op.column.name
619-
overlaps_with_parent = opname in op.parent.schema
620-
computed_column = column_alias.as_(opname, quoted=quoted)
625+
overlaps_with_parent = column_name in op.parent.schema
626+
computed_column = column_alias.as_(column_name, quoted=quoted)
621627

622628
selcols = []
623629

@@ -627,7 +633,9 @@ def visit_TableUnnest(
627633
# TODO: clean this up once WITH ORDINALITY is supported in DuckDB
628634
# no need for struct_extract once that's upstream
629635
column = self.f.list_zip(column, self.f.range(self.f.len(column)))
630-
extract = self.f.struct_extract(column_alias, 1).as_(opname, quoted=quoted)
636+
extract = self.f.struct_extract(column_alias, 1).as_(
637+
column_name, quoted=quoted
638+
)
631639

632640
if overlaps_with_parent:
633641
replace = sge.Column(this=sge.Star(replace=[extract]), table=table)

‎ibis/backends/sql/compilers/postgres.py

+10-4
Original file line numberDiff line numberDiff line change
@@ -718,18 +718,24 @@ def visit_Hash(self, op, *, arg):
718718
)
719719

720720
def visit_TableUnnest(
721-
self, op, *, parent, column, offset: str | None, keep_empty: bool
721+
self,
722+
op,
723+
*,
724+
parent,
725+
column,
726+
column_name: str,
727+
offset: str | None,
728+
keep_empty: bool,
722729
):
723730
quoted = self.quoted
724731

725732
column_alias = sg.to_identifier(gen_name("table_unnest_column"), quoted=quoted)
726733

727734
parent_alias = parent.alias_or_name
728735

729-
opname = op.column.name
730736
parent_schema = op.parent.schema
731-
overlaps_with_parent = opname in parent_schema
732-
computed_column = column_alias.as_(opname, quoted=quoted)
737+
overlaps_with_parent = column_name in parent_schema
738+
computed_column = column_alias.as_(column_name, quoted=quoted)
733739

734740
selcols = []
735741

‎ibis/backends/sql/compilers/pyspark.py

+10-4
Original file line numberDiff line numberDiff line change
@@ -500,16 +500,22 @@ def visit_HexDigest(self, op, *, arg, how):
500500
raise NotImplementedError(f"No available hashing function for {how}")
501501

502502
def visit_TableUnnest(
503-
self, op, *, parent, column, offset: str | None, keep_empty: bool
503+
self,
504+
op,
505+
*,
506+
parent,
507+
column,
508+
column_name: str,
509+
offset: str | None,
510+
keep_empty: bool,
504511
):
505512
quoted = self.quoted
506513

507514
column_alias = sg.to_identifier(gen_name("table_unnest_column"), quoted=quoted)
508515

509-
opname = op.column.name
510516
parent_schema = op.parent.schema
511-
overlaps_with_parent = opname in parent_schema
512-
computed_column = column_alias.as_(opname, quoted=quoted)
517+
overlaps_with_parent = column_name in parent_schema
518+
computed_column = column_alias.as_(column_name, quoted=quoted)
513519

514520
parent_alias = parent.alias_or_name
515521

‎ibis/backends/sql/compilers/snowflake.py

+11-6
Original file line numberDiff line numberDiff line change
@@ -810,7 +810,14 @@ def visit_DropColumns(self, op, *, parent, columns_to_drop):
810810
return sg.select(column).from_(parent)
811811

812812
def visit_TableUnnest(
813-
self, op, *, parent, column, offset: str | None, keep_empty: bool
813+
self,
814+
op,
815+
*,
816+
parent,
817+
column,
818+
column_name: str,
819+
offset: str | None,
820+
keep_empty: bool,
814821
):
815822
quoted = self.quoted
816823

@@ -825,12 +832,10 @@ def visit_TableUnnest(
825832

826833
selcols = []
827834

828-
opcol = op.column
829-
opname = opcol.name
830-
overlaps_with_parent = opname in op.parent.schema
835+
overlaps_with_parent = column_name in op.parent.schema
831836
computed_column = self.cast(
832-
self.f.nullif(column_alias, null_sentinel), opcol.dtype.value_type
833-
).as_(opname, quoted=quoted)
837+
self.f.nullif(column_alias, null_sentinel), op.column.dtype.value_type
838+
).as_(column_name, quoted=quoted)
834839

835840
if overlaps_with_parent:
836841
selcols.append(

‎ibis/backends/sql/compilers/trino.py

+10-4
Original file line numberDiff line numberDiff line change
@@ -546,16 +546,22 @@ def visit_ToJSONArray(self, op, *, arg):
546546
)
547547

548548
def visit_TableUnnest(
549-
self, op, *, parent, column, offset: str | None, keep_empty: bool
549+
self,
550+
op,
551+
*,
552+
parent,
553+
column,
554+
column_name: str,
555+
offset: str | None,
556+
keep_empty: bool,
550557
):
551558
quoted = self.quoted
552559

553560
column_alias = sg.to_identifier(gen_name("table_unnest_column"), quoted=quoted)
554561

555-
opname = op.column.name
556562
parent_schema = op.parent.schema
557-
overlaps_with_parent = opname in parent_schema
558-
computed_column = column_alias.as_(opname, quoted=quoted)
563+
overlaps_with_parent = column_name in parent_schema
564+
computed_column = column_alias.as_(column_name, quoted=quoted)
559565

560566
parent_alias_or_name = parent.alias_or_name
561567

‎ibis/backends/sql/rewrites.py

+8-1
Original file line numberDiff line numberDiff line change
@@ -165,7 +165,7 @@ def fill_null_to_select(_, **kwargs):
165165
for name in _.parent.schema.names:
166166
col = ops.Field(_.parent, name)
167167
if (value := mapping.get(name)) is not None:
168-
col = ops.Alias(ops.Coalesce((col, value)), name)
168+
col = ops.Coalesce((col, value))
169169
selections[name] = col
170170

171171
return Select(_.parent, selections=selections)
@@ -206,6 +206,12 @@ def first_to_firstvalue(_, **kwargs):
206206
return _.copy(func=klass(_.func.arg))
207207

208208

209+
@replace(p.Alias)
210+
def remove_aliases(_, **kwargs):
211+
"""Remove all remaining aliases, they're not needed for remaining compilation."""
212+
return _.arg
213+
214+
209215
def complexity(node):
210216
"""Assign a complexity score to a node.
211217
@@ -372,6 +378,7 @@ def sqlize(
372378
context = {"params": params}
373379
result = node.replace(
374380
replace_parameter
381+
| remove_aliases
375382
| project_to_select
376383
| filter_to_select
377384
| sort_to_select

‎ibis/backends/tests/sql/snapshots/test_compiler/test_subquery_where_location/decompiled.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
},
1212
)
1313
param = ibis.param("timestamp")
14-
f = alltypes.filter((alltypes.timestamp_col < param.name("my_param")))
14+
f = alltypes.filter((alltypes.timestamp_col < param))
1515
agg = f.aggregate([f.float_col.sum().name("foo")], by=[f.string_col])
1616

1717
result = agg.foo.count()

‎ibis/backends/tests/sql/snapshots/test_sql/test_binop_with_alias_still_parenthesized/out.sql

+5
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
SELECT
2+
(
3+
"t0"."a" + "t0"."b"
4+
) * "t0"."c" AS "x"
5+
FROM "t" AS "t0"

‎ibis/backends/tests/sql/test_compiler.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -196,7 +196,7 @@ def test_subquery_where_location(snapshot):
196196
],
197197
name="alltypes",
198198
)
199-
param = ibis.param("timestamp").name("my_param")
199+
param = ibis.param("timestamp")
200200
expr = (
201201
t[["float_col", "timestamp_col", "int_col", "string_col"]][
202202
lambda t: t.timestamp_col < param

‎ibis/backends/tests/sql/test_sql.py

+6
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,12 @@ def test_binop_parens(snapshot, opname, dtype, associative):
143143
snapshot.assert_match(combined, "out.sql")
144144

145145

146+
def test_binop_with_alias_still_parenthesized(snapshot):
147+
t = ibis.table({"a": "int", "b": "int", "c": "int"}, name="t")
148+
sql = to_sql(((t.a + t.b).name("d") * t.c).name("x"))
149+
snapshot.assert_match(sql, "out.sql")
150+
151+
146152
@pytest.mark.parametrize(
147153
"expr_fn",
148154
[

‎ibis/expr/operations/relations.py

+4-7
Original file line numberDiff line numberDiff line change
@@ -498,6 +498,7 @@ class TableUnnest(Relation):
498498

499499
parent: Relation
500500
column: Value[dt.Array]
501+
column_name: str
501502
offset: typing.Union[str, None]
502503
keep_empty: bool
503504

@@ -507,15 +508,11 @@ def values(self):
507508

508509
@attribute
509510
def schema(self):
510-
column = self.column
511-
offset = self.offset
512-
513511
base = self.parent.schema.fields.copy()
512+
base[self.column_name] = self.column.dtype.value_type
514513

515-
base[column.name] = column.dtype.value_type
516-
517-
if offset is not None:
518-
base[offset] = dt.int64
514+
if self.offset is not None:
515+
base[self.offset] = dt.int64
519516

520517
return Schema(base)
521518

‎ibis/expr/types/relations.py

+5-1
Original file line numberDiff line numberDiff line change
@@ -4890,7 +4890,11 @@ def unnest(
48904890
"""
48914891
(column,) = self.bind(column)
48924892
return ops.TableUnnest(
4893-
parent=self, column=column, offset=offset, keep_empty=keep_empty
4893+
parent=self,
4894+
column=column,
4895+
column_name=column.get_name(),
4896+
offset=offset,
4897+
keep_empty=keep_empty,
48944898
).to_expr()
48954899

48964900

0 commit comments

Comments
 (0)