feat(polars): support is_in queries from uncorrelated tables

Author: cpcloud

ibis/backends/polars/compiler.py

@@ -477,9 +477,29 @@ def greatest(op, **kw):
 
 @translate.register(ops.InSubquery)
 def in_column(op, **kw):
-    value = translate(op.value, **kw)
+    if op.value.dtype.is_struct():
+        raise com.UnsupportedOperationError(
+            "polars doesn't support contains with struct elements"
+        )
+
     needle = translate(op.needle, **kw)
-    return needle.is_in(value)
+    value = translate(op.value, **kw)
+    (rel,) = op.value.relations
+    # The `collect` triggers computation, but there appears to be no way to
+    # spell this operation in a polars-native way that's
+    #
+    # 1. not deprecated
+    # 2. operates only using pl.Expr objects and methods
+    #
+    # In other words, we need to either rearchitect the polars compiler to
+    # operate only with DataFrames, or compute first. I chose computing first
+    # since it is less effort and we don't know how impactful it is.
+    value = translate(rel, **kw).select(value).collect().to_series()
+
+    return needle.map_batches(
+        lambda needle, value=value: needle.is_in(value),
+        return_dtype=pl.Boolean(),
+    )
 
 
 @translate.register(ops.InValues)

ibis/backends/tests/test_generic.py

@@ -1128,7 +1128,32 @@ def test_typeof(con):
     assert result is not None
 
 
-@pytest.mark.notimpl(["polars"], reason="incorrect answer")
+@pytest.mark.notyet(["impala"], reason="can't find table in subquery")
+@pytest.mark.notimpl(["datafusion", "druid"])
+@pytest.mark.xfail_version(pyspark=["pyspark<3.5"])
+@pytest.mark.notyet(["exasol"], raises=ExaQueryError, reason="not supported by exasol")
+@pytest.mark.notyet(
+    ["risingwave"],
+    raises=PsycoPg2InternalError,
+    reason="https://github.com/risingwavelabs/risingwave/issues/1343",
+)
+@pytest.mark.notyet(
+    ["mssql"],
+    raises=PyODBCProgrammingError,
+    reason="naked IN queries are not supported",
+)
+def test_isin_uncorrelated_simple(con):
+    u1 = ibis.memtable({"id": [1, 2, 3]})
+    a = ibis.memtable({"id": [1, 2]})
+
+    u2 = u1.mutate(in_a=u1["id"].isin(a["id"]))
+    final = u2.order_by("id")
+
+    result = con.to_pyarrow(final)
+    expected = pa.table({"id": [1, 2, 3], "in_a": [True, True, False]})
+    assert result.equals(expected)
+
+
 @pytest.mark.notyet(["impala"], reason="can't find table in subquery")
 @pytest.mark.notimpl(["datafusion", "druid"])
 @pytest.mark.xfail_version(pyspark=["pyspark<3.5"])
@@ -1161,7 +1186,6 @@ def test_isin_uncorrelated(
     backend.assert_series_equal(result, expected)
 
 
-@pytest.mark.notimpl(["polars"], reason="incorrect answer")
 def test_isin_uncorrelated_filter(
     backend, batting, awards_players, batting_df, awards_players_df
 ):

ibis/backends/tests/test_struct.py

@@ -12,7 +12,6 @@
 from ibis.backends.tests.conftest import NO_STRUCT_SUPPORT_MARKS
 from ibis.backends.tests.errors import (
     DatabricksServerOperationError,
-    PolarsColumnNotFoundError,
     PsycoPg2InternalError,
     PsycoPg2ProgrammingError,
     PsycoPgSyntaxError,
@@ -21,7 +20,7 @@
     PyAthenaOperationalError,
     PySparkAnalysisException,
 )
-from ibis.common.exceptions import IbisError
+from ibis.common.exceptions import IbisError, UnsupportedOperationError
 
 np = pytest.importorskip("numpy")
 pd = pytest.importorskip("pandas")
@@ -240,8 +239,8 @@ def test_keyword_fields(con, nullable):
 )
 @pytest.mark.notyet(
     ["polars"],
-    raises=PolarsColumnNotFoundError,
-    reason="doesn't seem to support IN-style subqueries on structs",
+    raises=UnsupportedOperationError,
+    reason="doesn't support IN-style subqueries on structs",
 )
 @pytest.mark.xfail_version(
     pyspark=["pyspark<3.5"],