Определите, содержит ли группа значение в Polars ⇐ Python

[Anonymous]

Мне хотелось бы вернуть строки набора данных, в которых группы удовлетворяют двум условиям. То есть

• Я хочу знать, превышает ли количество уникальных значений в группе 1, и
• Я хочу знать, содержит ли группа определенное значение.

Я зацикливаюсь на последнем!
Предположим, у меня есть этот фрейм данных:

Код: Выделить всё

import polars as pl

df = pl.from_repr("""
┌──────┬──────────┬─────────┐
│ item ┆ location ┆ store   │
│ ---  ┆ ---      ┆ ---     │
│ i64  ┆ str      ┆ str     │
╞══════╪══════════╪═════════╡
│ 0    ┆ new york ┆ store 1 │
│ 0    ┆ boston   ┆ store 1 │
│ 1    ┆ boston   ┆ store 1 │
│ 1    ┆ boston   ┆ store 2 │
│ 0    ┆ ohio     ┆ store 1 │
│ 0    ┆ ohio     ┆ store 3 │
└──────┴──────────┴─────────┘
""")

Я могу легко проверить номер 1, используя фильтр, n_unique() и более следующим образом:

Код: Выделить всё

df.filter(
(pl.col("store").n_unique() > 1).over("item", "location")
)

Код: Выделить всё

┌──────┬──────────┬─────────┐
│ item ┆ location ┆ store   │
│ ---  ┆ ---      ┆ ---     │
│ i64  ┆ str      ┆ str     │
╞══════╪══════════╪═════════╡
│ 1    ┆ boston   ┆ store 1 │
│ 1    ┆ boston   ┆ store 2 │
│ 0    ┆ ohio     ┆ store 1 │
│ 0    ┆ ohio     ┆ store 3 │
└──────┴──────────┴─────────┘

Однако я хотел бы добавить в фильтр возврат строк, в которых соответствующая группа содержит «магазин 3».
Я попробовал это (и его варианты):

Код: Выделить всё

df.filter(
(pl.col("store").n_unique() > 1).over("item", "location")
& (pl.col("store").is_in(["store 3"]))
)

Код: Выделить всё

┌──────┬──────────┬─────────┐
│ item ┆ location ┆ store   │
│ ---  ┆ ---      ┆ ---     │
│ i64  ┆ str      ┆ str     │
╞══════╪══════════╪═════════╡
│ 0    ┆ ohio     ┆ store 3 │
└──────┴──────────┴─────────┘

Но это возвращает только одну строку. И на самом деле результат для меня имеет некоторый смысл, но я надеялся вернуть всю группу. Например,

Код: Выделить всё

┌──────┬──────────┬─────────┐
│ item ┆ location ┆ store   │
│ ---  ┆ ---      ┆ ---     │
│ i64  ┆ str      ┆ str     │
╞══════╪══════════╪═════════╡
│ 0    ┆ ohio     ┆ store 1 │
│ 0    ┆ ohio     ┆ store 3 │
└──────┴──────────┴─────────┘

Я совершенно уверен, что упускаю что-то очевидное. Любая помощь приветствуется. Без сомнения, до сих пор люблю Polars!

Подробнее здесь: https://stackoverflow.com/questions/770 ... -in-polars