diff --git a/ibis/backends/pandas/execution/arrays.py b/ibis/backends/pandas/execution/arrays.py index 8b23a5c93f43..234d4c49936a 100644 --- a/ibis/backends/pandas/execution/arrays.py +++ b/ibis/backends/pandas/execution/arrays.py @@ -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. diff --git a/ibis/backends/tests/test_map.py b/ibis/backends/tests/test_map.py index 45e6dafaad91..30d692df580d 100644 --- a/ibis/backends/tests/test_map.py +++ b/ibis/backends/tests/test_map.py @@ -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") diff --git a/ibis/expr/operations/logical.py b/ibis/expr/operations/logical.py index 4aaa9f347fe9..a0241214d467 100644 --- a/ibis/expr/operations/logical.py +++ b/ibis/expr/operations/logical.py @@ -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 diff --git a/ibis/expr/types/generic.py b/ibis/expr/types/generic.py index 5ed84705ab79..b7aa92671c24 100644 --- a/ibis/expr/types/generic.py +++ b/ibis/expr/types/generic.py @@ -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`. diff --git a/ibis/tests/expr/test_value_exprs.py b/ibis/tests/expr/test_value_exprs.py index 51bfe572a325..851f7c2d4731 100644 --- a/ibis/tests/expr/test_value_exprs.py +++ b/ibis/tests/expr/test_value_exprs.py @@ -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")], 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():