Mutating in a loop with isin using expressions as value doesn't work as expected

[jklaise]

Given a set of tables a, b, ... with the same column name id, my goal is to create a new table with a column id containing all possible id values and for each table a boolean column in_a, in_b, ... containing true if the specific table also has the specific id value. Sample code:

import ibis
import polars as pl

u = pl.DataFrame({"id": [1, 2, 3]})  # all possible ids
a = pl.DataFrame({"id": [1, 2]})
b = pl.DataFrame({"id": [2, 3]})

con = ibis.polars.connect({"u": u, "a": a, "b": b})

# desired output:
# id,   in_a,   in_b
# 1,    true,   false
# 2,    true,   true
# 3,    false,  true

u = con.table("u")
t_names = ["a", "b"]
tables = [con.table(t) for t in t_names]

for t, name in zip(tables, t_names):
    u = u.mutate(**{f"in_{name}": u["id"].isin(t["id"])})

print(u.to_polars())
# actual output:
# id,   in_a,   in_b
# 1,    true,   true
# 2,    true,   true
# 3,    true,   true

However, changing the argument to isin to be static, e.g.

u = u.mutate(**{f"in_{name}": u["id"].isin(t["id"]).to_polars()})

gives the right answer. Wondering if this is expected behaviour or a bug?

[jklaise]

I think I've narrowed down the issue to dynamic lookup using expressions inside isin not working on the polars backend. For the duckdb backend this works as expected:

import ibis

u = ibis.memtable({"id": [1, 2, 3]})
a = ibis.memtable({"id": [1, 2]})
b = ibis.memtable({"id": [2, 3]})

tables = {"a": a, "b": b}
for name, t in tables.items():
    u = u.mutate(**{f"in_{name}": u["id"].isin(t["id"])})
print(u)

┏━━━━━━━┳━━━━━━━━━┳━━━━━━━━━┓
┃ id    ┃ in_a    ┃ in_b    ┃
┡━━━━━━━╇━━━━━━━━━╇━━━━━━━━━┩
│ int64 │ boolean │ boolean │
├───────┼─────────┼─────────┤
│     1 │ True    │ False   │
│     2 │ True    │ True    │
│     3 │ False   │ True    │
└───────┴─────────┴─────────┘

But setting backed to polars results in the wrong answer:

import ibis

ibis.set_backend("polars")

u = ibis.memtable({"id": [1, 2, 3]})
a = ibis.memtable({"id": [1, 2]})
b = ibis.memtable({"id": [2, 3]})

tables = {"a": a, "b": b}
for name, t in tables.items():
    u = u.mutate(**{f"in_{name}": u["id"].isin(t["id"])})

┏━━━━━━━┳━━━━━━━━━┳━━━━━━━━━┓
┃ id    ┃ in_a    ┃ in_b    ┃
┡━━━━━━━╇━━━━━━━━━╇━━━━━━━━━┩
│ int64 │ boolean │ boolean │
├───────┼─────────┼─────────┤
│     1 │ True    │ True    │
│     2 │ True    │ True    │
│     3 │ True    │ True    │
└───────┴─────────┴─────────┘

[jklaise]

The workaround with a join (+ mutate after) works:

upl = pl.DataFrame({"id": [1, 2, 3]})  # all posisble ids
apl = pl.DataFrame({"id": [1, 2]})
bpl = pl.DataFrame({"id": [2,3]})

u, a, b = ibis.memtable(upl), ibis.memtable(apl), ibis.memtable(bpl)
tables = {"a": a, "b": b}

for name, t in tables.items():
    u = u.join(t, u.id==t.id, how="left", rname=f"in_{name}").mutate(**{f"in_{name}": _[f"in_{name}"].notnull()})
u

┏━━━━━━━┳━━━━━━━━━┳━━━━━━━━━┓
┃ id    ┃ in_a    ┃ in_b    ┃
┡━━━━━━━╇━━━━━━━━━╇━━━━━━━━━┩
│ int64 │ boolean │ boolean │
├───────┼─────────┼─────────┤
│     1 │ True    │ False   │
│     2 │ True    │ True    │
│     3 │ False   │ True    │
└───────┴─────────┴─────────┘

[cpcloud]

What is the non-deprecated, pure polars solution to this? We can try to use that inside Ibis.

[jklaise]

I'm not sure there is one, reading pola-rs/polars#22149, the reason seems tangential, e.g. to disambiguate when looking up things in list rows. The warning given when using polars is_in looking up values in the series of another dataframe suggests to materialize the lookup values to a Python list. But anyway, in my example upl.with_columns(in_a=pl.col("id").is_in(apl["id"])), the apl["id"] is not an expression but a materialized series already.

This thread might be relevant - the suggestions seem to range from using lazyframes to using joins instead of lookups which should be faster in any case: pola-rs/polars#4670

[cpcloud]

Not entirely following why this is an anti pattern, but 🤷🏻.

If you're immediately using the boolean column as a filter, then yes, a semi or anti join is what you probably should be using, though most query engines recognize the pattern and convert it to one of those for you as part of query optimization.

If you need the boolean column for some non-filtration purpose, then it seems like your best bet is to use joins for now.

I tried an approach that gets this working, but it uses a combination of map_batches and materializing one side of the is_in, so it's pretty much a performance hellscape.

[cpcloud]

Zooming out a bit, architecturally Ibis assumes that all polars operations can be done with pl.Expr, which I guess is technically true, but it prevents us from easily substituting in equivalent expressions for tables and then pulling a column from that.

Though ... now that I think about it, maybe there's a way we can rewrite InSubquery operations in Ibis into something Polars won't consider an anti-pattern 🤔