Skip to content

Commit

Permalink
fix: fix table.mutate with deferred named expressions
Browse files Browse the repository at this point in the history
  • Loading branch information
jcrist authored and cpcloud committed Oct 15, 2022
1 parent 1b32856 commit 5877d0b
Show file tree
Hide file tree
Showing 3 changed files with 42 additions and 13 deletions.
1 change: 0 additions & 1 deletion ibis/expr/analysis.py
Original file line number Diff line number Diff line change
Expand Up @@ -229,7 +229,6 @@ def get_mutation_exprs(exprs: list[ir.Expr], table: ir.Table) -> list[ir.Expr |
# name does not exist in the original table.
# Given these two data structures, we can compute the mutation node exprs
# based on whether any columns are being overwritten.
# TODO issue #2649
overwriting_cols_to_expr: dict[str, ir.Expr | None] = {}
non_overwriting_exprs: list[ir.Expr] = []
table_schema = table.schema()
Expand Down
18 changes: 11 additions & 7 deletions ibis/expr/types/relations.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
import collections
import functools
import itertools
import operator
import sys
import warnings
from typing import (
Expand Down Expand Up @@ -590,16 +589,21 @@ def mutate(
import ibis.expr.analysis as an
import ibis.expr.rules as rlz

exprs = [] if exprs is None else util.promote_list(exprs)
for name, expr in sorted(mutations.items(), key=operator.itemgetter(0)):
def ensure_expr(expr):
# This is different than self._ensure_expr, since we don't want to
# treat `str` or `int` values as column indices
if util.is_function(expr):
value = expr(self)
return expr(self)
elif isinstance(expr, Deferred):
value = expr.resolve(self)
return expr.resolve(self)
else:
value = rlz.any(expr).to_expr()
exprs.append(value.name(name))
return rlz.any(expr).to_expr()

exprs = [] if exprs is None else util.promote_list(exprs)
exprs = [ensure_expr(expr) for expr in exprs]
exprs.extend(
ensure_expr(expr).name(name) for name, expr in sorted(mutations.items())
)
mutation_exprs = an.get_mutation_exprs(exprs, self)
return self.select(mutation_exprs)

Expand Down
36 changes: 31 additions & 5 deletions ibis/tests/expr/test_table.py
Original file line number Diff line number Diff line change
Expand Up @@ -238,11 +238,37 @@ def test_projection_array_expr(table):


def test_mutate(table):
one = table.f * 2
foo = (table.a + table.b).name('foo')

expr = table.mutate(foo, one=one, two=2)
expected = table[table, foo, one.name('one'), ibis.literal(2).name('two')]
expr = table.mutate(
[
(table.a + 1).name("x1"),
table.b.sum().name("x2"),
(_.a + 2).name("x3"),
lambda _: (_.a + 3).name("x4"),
4,
"five",
],
kw1=(table.a + 6),
kw2=table.b.sum(),
kw3=(_.a + 7),
kw4=lambda _: (_.a + 8),
kw5=9,
kw6="ten",
)
expected = table[
table,
(table.a + 1).name("x1"),
table.b.sum().name("x2"),
(table.a + 2).name("x3"),
(table.a + 3).name("x4"),
ibis.literal(4).name("4"),
ibis.literal("five").name("'five'"),
(table.a + 6).name("kw1"),
table.b.sum().name("kw2"),
(table.a + 7).name("kw3"),
(table.a + 8).name("kw4"),
ibis.literal(9).name("kw5"),
ibis.literal("ten").name("kw6"),
]
assert_equal(expr, expected)


Expand Down

0 comments on commit 5877d0b

Please sign in to comment.