Условная дедупликация в полярах

У меня есть набор данных, из которого я пытаюсь удалить повторяющиеся записи. Lazyframe, с которым я работаю, структурирован следующим образом:

df = pl.from_repr("""
┌──────┬────────────┬──────────────────┬───────┐
│ id   ┆ title      ┆ type             ┆ type2 │
│ ---  ┆ ---        ┆ ---              ┆ ---   │
│ i64  ┆ str        ┆ str              ┆ i64   │
╞══════╪════════════╪══════════════════╪═══════╡
│ 1001 ┆ Research A ┆ journal article  ┆ 35    │
│ 1002 ┆ Research B ┆ book chapter     ┆ 41    │
│ 1003 ┆ Research C ┆ journal article  ┆ 35    │
│ 1004 ┆ Research D ┆ conference paper ┆ 42    │
│ 1001 ┆ Research E ┆ journal article  ┆ 35    │
│ 1002 ┆ Research F ┆ journal article  ┆ 41    │
│ 1003 ┆ Research G ┆                  ┆ 41    │
│ 1002 ┆ Research I ┆ book chapter     ┆ 41    │
│ 1003 ┆ Research J ┆ journal article  ┆ 35    │
└──────┴────────────┴──────────────────┴───────┘
""")

Я хочу удалить записи с одинаковым id, но на самом деле бывают разные случаи:

  1. Дубликаты имеют одинаковый type (например, 1001): сохраните первый.
  2. Дубликаты имеют разные type: отбрасывайте те, которые имеют пустую строку ("") в качестве типа, а затем сохраняйте только те записи, которые соответствуют следующим парам type и type2:
dict_df = pl.DataFrame({
    "type": ['journal article', 'book chapter', 'book chapter'],
    "type2": [35, 41, 42]
})

Ожидаемый результат

идентификатор [i64] заголовок[строка] тип[строка] тип2[i64] 1001 Исследование А журнальная статья 35 1002 Исследование Б глава книги 41 1003 Исследования С журнальная статья 35 1004 Исследования Д доклад конференции 42
  • 1001: тот же тип, оставьте первый
  • 1002: другой тип, сохраните первое вхождение записей с парой {'book Chapter': 41}
  • 1003: другой тип, отбросить записи с пустым типом и сохранить первое вхождение
  • 1004: не дубликат

Я пробовал много вещей, в основном используя выражение pl.when(), но не смог найти способ фильтровать группы.

(
    df
    .sort('type', descending=True)
    .group_by("id")
    .agg([
        pl.when(pl.col("type").n_unique() > 1)
        .then( ... )
        .otherwise(pl.all().first())
    ])
)

🤔 А знаете ли вы, что...
Python популярен в машинном обучении и искусственном интеллекте.


4
54
2

Ответы:

Вероятно, вы могли бы применить свою логику в контексте GroupBy, но также вы можете сначала предварительно отфильтровать DataFrame, чтобы у вас были только те строки, которые вы хотите сохранить:

(
    df
    .join(dict_df.with_row_index(), on=["type","type2"], how = "left")
    .filter(
        (pl.col.type.n_unique().over("id") == 1) |
        ((pl.col.type != "") & (pl.col.index.is_not_null()))
    )
    .drop("index")
    .unique("id", keep = "first", maintain_order=True)
    # alternatively,
    # .group_by("id").first()
)

┌──────┬────────────┬──────────────────┬───────┐
│ id   ┆ title      ┆ type             ┆ type2 │
│ ---  ┆ ---        ┆ ---              ┆ ---   │
│ i64  ┆ str        ┆ str              ┆ i64   │
╞══════╪════════════╪══════════════════╪═══════╡
│ 1001 ┆ Research A ┆ journal article  ┆ 35    │
│ 1003 ┆ Research C ┆ journal article  ┆ 35    │
│ 1004 ┆ Research D ┆ conference paper ┆ 42    │
│ 1002 ┆ Research F ┆ journal article  ┆ 35    │
└──────┴────────────┴──────────────────┴───────┘

Решено
  1. Создайте таблицу предпочтений, указав, сколько вы хотите каждой комбинации:
preference = pl.DataFrame({
    "type": ["journal article", "book chapter", "book chapter"],
    "iris_type": [35, 41, 42],
    "preference": [0, 1, 2]
})
  1. Соедините таблицу предпочтений с вашей таблицей данных:
joined = df.lazy().join(preference.lazy(), on=["type", "iris_type"], how = "left")
  1. Отсортируйте объединенную таблицу по предпочтениям, выберите первую в каждой группе и удалите столбец предпочтений:
out = (
    joined.sort("preference", descending=True, nulls_last=True)
        .group_by("id")
        .first()
        .drop("preference")
        .collect()
)