Skip to content

Commit

Permalink
chore: handle overwriting the unnested column
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud committed Jun 21, 2024
1 parent ff89731 commit ceddad0
Show file tree
Hide file tree
Showing 3 changed files with 21 additions and 6 deletions.
18 changes: 14 additions & 4 deletions ibis/backends/sql/compiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -1514,15 +1514,25 @@ def visit_Unsupported(self, op, **_):

def visit_TableUnnest(self, op, *, parent, column, offset: str | None):
quoted = self.quoted
parent_schema = op.parent.schema
parent_alias = parent.alias_or_name
column_name = op.column.name

column_alias = sg.to_identifier(column_name, quoted=quoted)
columns = [
sg.column(k, parent.alias_or_name, quoted=quoted)
for k in op.parent.schema.keys()
sg.column(k, parent_alias, quoted=quoted)
if k != column_name
else column_alias
for k in parent_schema.keys()
]
if column_name not in parent_schema:
columns.append(column_alias)

if offset is not None:
offset = sg.to_identifier(offset, quoted=quoted)
columns.append(offset)
column_alias = sg.to_identifier(op.column.name, quoted=self.quoted)
expr = sg.select(*columns, column_alias)

expr = sg.select(*columns)
alias = sge.TableAlias(columns=[column_alias])
return expr.from_(parent).join(
sge.Unnest(expressions=[column], alias=alias, offset=offset),
Expand Down
2 changes: 1 addition & 1 deletion ibis/backends/tests/test_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -1392,6 +1392,6 @@ def test_zip_unnest_lift(con):

def test_table_unnest(backend):
t = backend.array_types
expr = t.unnest(t.x.name("unnested_x"), offset="i")
expr = t.unnest(t.x)
result = expr.execute()
assert len(result)
7 changes: 6 additions & 1 deletion ibis/expr/operations/relations.py
Original file line number Diff line number Diff line change
Expand Up @@ -480,7 +480,12 @@ def schema(self):
column = self.column
offset = self.offset

base = self.parent.schema | Schema({column.name: column.dtype.value_type})
base = Schema.from_tuples(
(name, dtype.value_type) if name == column.name else (name, dtype)
for name, dtype in self.parent.schema.items()
)
if column.name not in base:
base |= Schema({column.name: column.dtype.value_type})
if offset is not None:
base |= Schema({offset: dt.int64})
return base
Expand Down

0 comments on commit ceddad0

Please sign in to comment.