Polars中基于多规则的数据集条件去重问题求助

方法一：使用 map_groups 自定义分组处理 #

这种方法逻辑直观，便于理解分组内的处理规则：

步骤说明 #

1. 预先构建有效 (type, type2) 配对集合，用于快速校验。
2. 过滤掉所有 type 为空字符串的行。
3. 按 id 分组，对每个分组应用自定义逻辑：
   • 若分组内所有 type 相同，保留第一行。
   • 若分组内 type 不同，过滤出符合有效配对的行后保留第一行。

代码实现 #

import polars as pl

# 原始数据集
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    │
└──────┴────────────┴──────────────────┴───────┘
""")

# 有效配对规则
dict_df = pl.DataFrame({
    "type": ['journal article', 'book chapter', 'book chapter'],
    "type2": [35, 41, 42]
})
valid_pairs = set(tuple(row) for row in dict_df.rows())

# 分组处理函数
def process_group(group: pl.DataFrame) -> pl.DataFrame:
    if group["type"].n_unique() == 1:
        return group.head(1)
    else:
        valid_rows = group.filter(pl.struct("type", "type2").is_in(valid_pairs))
        return valid_rows.head(1)

# 执行去重（支持 LazyFrame）
result = (
    df.lazy()
    .filter(pl.col("type") != "")
    .group_by("id")
    .map_groups(process_group)
    .collect()
)

print(result)

方法二：使用窗口函数（高性能推荐） #

这种方法采用向量化操作，适合处理大规模数据集，性能优于 map_groups：

步骤说明 #

1. 构建有效配对集合。
2. 通过窗口函数标记每个分组是否存在多种 type、每行是否符合有效配对、以及行的原始顺序。
3. 为每行计算优先级：
   • type 为空的行优先级为 0（被过滤）。
   • 分组内 type 唯一的行优先级为 2。
   • 分组内 type 多样且符合有效配对的行优先级为 1。
4. 按优先级和原始顺序排序，保留每个分组的第一行。

代码实现 #

import polars as pl

# 原始数据集与有效配对规则同方法一

valid_pairs = set(tuple(row) for row in dict_df.rows())

result = (
    df.lazy()
    .with_columns(
        # 标记分组是否有多种type
        has_multiple_types=pl.col("type").n_unique().over("id") > 1,
        # 标记行是否符合有效配对
        is_valid=pl.struct("type", "type2").is_in(valid_pairs),
        # 记录分组内的原始行顺序
        row_num=pl.int_range(0, pl.count()).over("id")
    )
    .with_columns(
        # 计算优先级
        priority=pl.when(pl.col("type") == "").then(0)
                   .when(~pl.col("has_multiple_types")).then(2)
                   .when(pl.col("is_valid")).then(1)
                   .otherwise(0)
    )
    .filter(pl.col("priority") > 0)  # 过滤低优先级行
    .with_columns(
        # 按优先级和原始顺序排名
        rank=pl.rank(descending=True, by=["priority", "row_num"]).over("id")
    )
    .filter(pl.col("rank") == 1)  # 保留每个分组的第一行
    .drop(["has_multiple_types", "is_valid", "row_num", "priority", "rank"])
    .collect()
)

print(result)

输出结果 #

两种方法均会得到预期输出：

┌──────┬────────────┬──────────────────┬───────┐
│ 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    │
└──────┴────────────┴──────────────────┴───────┘

内容的提问来源于stack exchange，提问作者lzill