Pandas isin function in polars

Question

Once in a while I get to the point where I need to run the following line:

DF[‘is_flagged’] = DF[‘id’].isin(DF2[DF2[‘flag’]==1][‘id’])

Lately I started using polars, and I wonder how to convert it easily to polars.

Edit: For example:

df1 = polars.DataFrame._from_dict({     
'Animal_id': [1, 2, 3, 4, 5, 6, 7],     
'age': [4, 6, 3, 8, 3, 8, 9] })

df2 = polars.DataFrame._from_dict({     
'Animal_id': [1, 2, 3, 4, 5, 6, 7],     
'Animal_type': ['cat', 'dog', 'cat', 'cat', 'dog', 'dog', 'cat'] })

Output:

polars.DataFrame._from_dict({
'animal_id': [1, 2, 3, 4, 5, 6, 7],
'age': [4, 6, 3, 8, 3, 8, 9],
'is_dog': [0, 1, 0, 0, 1, 1, 0]})

Without using flag and then join

I tried to use the is_in() function but this didn’t worked.

What is your desired output? let's say we have: `df = polars.DataFrame._from_dict({ 'flag': [1, 1, 0, 0, 1, 1, 1], 'id': ['hello', 'hello', 'world', 'world', 'hello', 'hello', 'hello'] })` - what would your desired output look like? — CodeCop, Jan 16 '23 at 19:42

score 1 · Answer 1 · answered Jan 16 '23 at 20:48

How about

df1.with_column(
    df2.with_column(
        pl.col("Animal_type")
        .is_in(["dog"])
        .cast(int)
        .alias("is_dog")
    )["is_dog"]
)

This gives you:

Out[13]:
shape: (7, 3)
┌───────────┬─────┬────────┐
│ Animal_id ┆ age ┆ is_dog │
│ ---       ┆ --- ┆ ---    │
│ i64       ┆ i64 ┆ i64    │
╞═══════════╪═════╪════════╡
│ 1         ┆ 4   ┆ 0      │
│ 2         ┆ 6   ┆ 1      │
│ 3         ┆ 3   ┆ 0      │
│ 4         ┆ 8   ┆ 0      │
│ 5         ┆ 3   ┆ 1      │
│ 6         ┆ 8   ┆ 1      │
│ 7         ┆ 9   ┆ 0      │
└───────────┴─────┴────────┘

score 1 · Accepted Answer · answered Jan 16 '23 at 22:07

So to be honest I am not quite sure from your question, your pandas snippet and your example what your desired solution is, but here are my three takes.

import polars as pl

df1 = pl.DataFrame(
    {"Animal_id": [1, 2, 3, 4, 5, 6, 7], "age": [4, 6, 3, 8, 3, 8, 9]}
).lazy()

df2 = pl.DataFrame(
    {
        "Animal_id": [1, 2, 3, 4, 5, 6, 7],
        "Animal_type": ["cat", "dog", "cat", "cat", "dog", "dog", "cat"],
    }
).lazy()

1 Solution

So this one is a small adaptation of the solution of @ignoring_gravity. So the assumption in his solution is that the DataFrames have the same length and the Animal_id matches in both tables. If that was your goal I want to give you another solution because by subsetting (["is_dog"]) you lose the possibility to use the lazy api.

df1.with_context(
    df2.select(pl.col("Animal_type").is_in(["dog"]).cast(int).alias("is_dog"))
).select(["Animal_id", "age", "is_dog"]).collect()

2 Solution

So in case you want something more similar to your pandas snippet and because you wrote you don't want to have a join.

df1.with_context(
    df2.filter(pl.col("Animal_type") == "dog").select(
        pl.col("Animal_id").alias("valid_ids")
    )
).with_columns(
    [pl.col("Animal_id").is_in(pl.col("valid_ids")).cast(int).alias("is_dog")]
).collect()

3 Solution

So this would be the solution with a join. In my opinion the best solution regarding your example and example output, but maybe there are other reasons that speak against a join, which aren't apparent from your example.

df1.join(
    df2.select(
        ["Animal_id", pl.col("Animal_type").is_in(["dog"]).cast(int).alias("is_dog")]
    ),
    on=["Animal_id"],
).collect()

Pandas isin function in polars

2 Answers2

1 Solution

2 Solution

3 Solution