基于窗口函数(mapping_strategy='join')按条件选首值并优化集合差计算

需要在使用mapping_strategy='join'的窗口函数（over）中，根据另一列的筛选条件选取某列的首个列表值，再用该值对两个列表列执行**集合差（SET DIFFERENCE）**操作。

示例数据 #

df = pl.DataFrame({
    'key': ['id1','id1', 'id1'],
    'list_main': [['a','b','c','d','e']]*3,
    "list_per_item":[['a','b'], ['a','b','c','d'], ['a','b','c']],
    'found': [False, False, True]
})
print(df)

输出：

shape: (3, 4)
┌─────┬───────────────────┬───────────────────┬───────┐
│ key ┆ list_main         ┆ list_per_item     ┆ found │
│ --- ┆ ---               ┆ ---               ┆ ---   │
│ str ┆ list[str]         ┆ list[str]         ┆ bool  │
╞═════╪═══════════════════╪═══════════════════╪═══════╡
│ id1 ┆ ["a", "b", … "e"] ┆ ["a", "b"]        ┆ false │
│ id1 ┆ ["a", "b", … "e"] ┆ ["a", "b", … "d"] ┆ false │
│ id1 ┆ ["a", "b", … "e"] ┆ ["a", "b", "c"]   ┆ true  │
└─────┴───────────────────┴───────────────────┴───────┘

现有实现 #

通过先按found降序排序，再用窗口函数取首值计算集合差：

value = pl.col('list_per_item').first().over('key', mapping_strategy='join')
df.sort('found', descending=True).with_columns(
    pl.col('list_per_item').list.set_difference(value),
    pl.col('list_main').list.set_difference(value),
    value.alias('list_to_subtract') # 实际代码中可删除此列
)

输出：

shape: (3, 5)
┌─────┬────────────┬───────────────┬───────┬──────────────────┐
│ key ┆ list_main  ┆ list_per_item ┆ found ┆ list_to_subtract │
│ --- ┆ ---        ┆ ---           ┆ ---   ┆ ---              │
│ str ┆ list[str]  ┆ list[str]     ┆ bool  ┆ list[str]        │
╞═════╪════════════╪═══════════════╪═══════╪══════════════════╡
│ id1 ┆ ["e", "d"] ┆ []            ┆ true  ┆ ["a", "b", "c"]  │
│ id1 ┆ ["e", "d"] ┆ []            ┆ false ┆ ["a", "b", "c"]  │
│ id1 ┆ ["e", "d"] ┆ ["d"]         ┆ false ┆ ["a", "b", "c"]  │
└─────┴────────────┴───────────────┴───────┴──────────────────┘

遇到的问题 #

尝试过先筛选DataFrame再关联的方式，但在30万行数据场景下出现OOM（内存溢出），需要更高效的表达式实现方式。

无需提前排序，直接在窗口函数内结合条件筛选获取目标列表，避免排序带来的额外内存开销，同时保留mapping_strategy='join'的窗口关联逻辑：

df.with_columns(
    # 按key分组，筛选found=True的行，取首个list_per_item作为待减集合
    list_to_subtract=pl.col('list_per_item').filter(pl.col('found')).first().over('key', mapping_strategy='join'),
    # 计算list_main与待减集合的差集
    list_main=pl.col('list_main').list.set_difference(pl.col('list_to_subtract')),
    # 计算list_per_item与待减集合的差集
    list_per_item=pl.col('list_per_item').list.set_difference(pl.col('list_to_subtract'))
).drop('list_to_subtract') # 可选：删除中间辅助列

方案说明 #

• 直接在窗口函数内用filter(pl.col('found'))精准筛选目标行，再取first()获取待减列表，省去全局排序步骤，大幅降低内存占用
• 保持mapping_strategy='join'的窗口关联逻辑，适配大数据集的分组计算场景
• 差集计算直接基于生成的辅助列，逻辑简洁且执行效率更高

输出结果 #

shape: (3, 4)
┌─────┬────────────┬───────────────┬───────┐
│ key ┆ list_main  ┆ list_per_item ┆ found │
│ --- ┆ ---        ┆ ---           ┆ ---   │
│ str ┆ list[str]  ┆ list[str]     ┆ bool  │
╞═════╪════════════╪═══════════════╪═══════╡
│ id1 ┆ ["e", "d"] ┆ []            ┆ false │
│ id1 ┆ ["e", "d"] ┆ ["d"]         ┆ false │
│ id1 ┆ ["e", "d"] ┆ []            ┆ true  │
└─────┴────────────┴───────────────┴───────┘

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