Skip to content

Commit

Permalink
chore: apply code review suggestions
Browse files Browse the repository at this point in the history
  • Loading branch information
NickCrews committed Feb 7, 2024
1 parent f26de65 commit c045312
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 18 deletions.
28 changes: 12 additions & 16 deletions ibis/backends/tests/test_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -449,26 +449,24 @@ def test_array_slice(backend, start, stop):
],
)
@pytest.mark.parametrize(
"func,result_name",
"func",
[
(lambda x: x + 1, "ArrayMap(a, Add(x, 1))"),
(functools.partial(lambda x, y: x + y, y=1), "ArrayMap(a, Add(x, 1))"),
(ibis._ + 1, "ArrayMap(a, Add(_, 1))"),
lambda x: x + 1,
functools.partial(lambda x, y: x + y, y=1),
ibis._ + 1,
],
)
@pytest.mark.broken(
["risingwave"],
raises=AssertionError,
reason="TODO(Kexiang): seems a bug",
)
def test_array_map(con, input, output, func, result_name):
def test_array_map(con, input, output, func):
t = ibis.memtable(input, schema=ibis.schema(dict(a="!array<int8>")))
t = ibis.memtable(input, schema=ibis.schema(dict(a="!array<int8>")))
expected = pd.Series(output["a"])

result = t.a.map(func)
assert result.get_name() == result_name
expr = t.select(a=result)
expr = t.select(a=t.a.map(func))
result = con.execute(expr.a)
assert frozenset(map(tuple, result.values)) == frozenset(
map(tuple, expected.values)
Expand Down Expand Up @@ -517,20 +515,18 @@ def test_array_map(con, input, output, func, result_name):
],
)
@pytest.mark.parametrize(
("predicate,result_name"),
"predicate",
[
(lambda x: x > 1, "ArrayFilter(a, Greater(x, 1))"),
(functools.partial(lambda x, y: x > y, y=1), "ArrayFilter(a, Greater(x, 1))"),
(ibis._ > 1, "ArrayFilter(a, Greater(_, 1))"),
lambda x: x > 1,
functools.partial(lambda x, y: x > y, y=1),
ibis._ > 1,
],
)
def test_array_filter(con, input, output, predicate, result_name):
def test_array_filter(con, input, output, predicate):
t = ibis.memtable(input, schema=ibis.schema(dict(a="!array<int8>")))
expected = pd.Series(output["a"])

result = t.a.filter(predicate)
assert result.get_name() == result_name
expr = t.select(a=result)
expr = t.select(a=t.a.filter(predicate))
result = con.execute(expr.a)
assert frozenset(map(tuple, result.values)) == frozenset(
map(tuple, expected.values)
Expand Down
4 changes: 2 additions & 2 deletions ibis/expr/types/arrays.py
Original file line number Diff line number Diff line change
Expand Up @@ -444,7 +444,7 @@ def map(self, func: Deferred | Callable[[ir.Value], ir.Value]) -> ir.ArrayValue:
└────────────────────────┘
"""
if isinstance(func, Deferred):
name = "_"
name = str(func)
else:
name = next(iter(inspect.signature(func).parameters.keys()))
parameter = ops.Argument(
Expand Down Expand Up @@ -544,7 +544,7 @@ def filter(
└───────────────────────────────┘
"""
if isinstance(predicate, Deferred):
name = "_"
name = str(predicate)
else:
name = next(iter(inspect.signature(predicate).parameters.keys()))
parameter = ops.Argument(
Expand Down

0 comments on commit c045312

Please sign in to comment.