Skip to content

Commit

Permalink
refactor(ir): construct ArrayContains instead of Contains for `va…
Browse files Browse the repository at this point in the history
…lue.isin(array_value)`
  • Loading branch information
deepyaman authored and cpcloud committed Aug 7, 2023
1 parent 7cfce45 commit e826037
Show file tree
Hide file tree
Showing 5 changed files with 30 additions and 8 deletions.
5 changes: 5 additions & 0 deletions ibis/backends/pandas/execution/arrays.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,11 @@ def execute_array_index_scalar(op, data, index, **kwargs):
return None


@execute_node.register(ops.ArrayContains, np.ndarray, object)
def execute_node_contains_value_array(op, haystack, needle, **kwargs):
return needle in haystack


def _concat_iterables_to_series(*iters: Collection[Any]) -> pd.Series:
"""Concatenate two collections to create a Series.
Expand Down
3 changes: 1 addition & 2 deletions ibis/backends/tests/test_map.py
Original file line number Diff line number Diff line change
Expand Up @@ -74,9 +74,8 @@ def test_literal_map_values(con):
assert np.array_equal(result, ["a", "b"])


@pytest.mark.notimpl(["trino", "postgres"])
@pytest.mark.notimpl(["postgres"])
@pytest.mark.notyet(["snowflake"])
@pytest.mark.notyet(["duckdb"], reason="sqlalchemy warning")
def test_scalar_isin_literal_map_keys(con):
mapping = ibis.literal({"a": 1, "b": 2})
a = ibis.literal("a")
Expand Down
6 changes: 1 addition & 5 deletions ibis/expr/operations/logical.py
Original file line number Diff line number Diff line change
Expand Up @@ -133,11 +133,7 @@ def __init__(self, arg, lower_bound, upper_bound):
@public
class Contains(Value):
value: Value
options: Union[
VarTuple[Value],
Column[dt.Any],
Value[dt.Array],
]
options: Union[VarTuple[Value], Column[dt.Any]]

dtype = dt.boolean

Expand Down
7 changes: 6 additions & 1 deletion ibis/expr/types/generic.py
Original file line number Diff line number Diff line change
Expand Up @@ -367,7 +367,12 @@ def isin(self, values: Value | Sequence[Value]) -> ir.BooleanValue:
other_string_col string
Contains(string_col, other_string_col): Contains(...)
"""
return ops.Contains(self, values).to_expr()
from ibis.expr.types import ArrayValue

if isinstance(values, ArrayValue):
return values.contains(self)
else:
return ops.Contains(self, values).to_expr()

def notin(self, values: Value | Sequence[Value]) -> ir.BooleanValue:
"""Check whether this expression's values are not in `values`.
Expand Down
17 changes: 17 additions & 0 deletions ibis/tests/expr/test_value_exprs.py
Original file line number Diff line number Diff line change
Expand Up @@ -1144,11 +1144,28 @@ def test_scalar_isin_map_keys():
assert isinstance(expr, ir.BooleanScalar)


def test_column_isin_array():
# scalar case
t = ibis.table([("a", "string")], name="t")
expr = t.a.isin(ibis.array(["a", "b"]))
assert isinstance(expr, ir.BooleanColumn)
assert isinstance(expr.op(), ops.ArrayContains)
assert expr.op().shape.is_columnar()

# columnar case
t = ibis.table([("a", "string"), ("b", "array<string>")], name="t")
expr = t.a.isin(t.b)
assert isinstance(expr, ir.BooleanColumn)
assert isinstance(expr.op(), ops.ArrayContains)
assert expr.op().shape.is_columnar()


def test_column_isin_map_keys():
t = ibis.table([("a", "string")], name="t")
mapping = ibis.literal({"a": 1, "b": 2})
expr = t.a.isin(mapping.keys())
assert isinstance(expr, ir.BooleanColumn)
assert isinstance(expr.op(), ops.ArrayContains)


def test_map_get_with_compatible_value_smaller():
Expand Down

0 comments on commit e826037

Please sign in to comment.