如何在Polars的groupby中实现Pandas的min_count参数功能？

需求说明 #

需要在Polars中实现类似Pandas groupby(key).sum(min_count=N)的效果：按fruit字段分组后，仅保留组内元素数量至少为4的分组，并计算该组price的求和结果。示例输入DataFrame如下：

import polars as pl

df = pl.from_repr("""
┌───────┬───────┐
│ fruit ┆ price │
│ ---   ┆ ---   │
│ str   ┆ i64   │
╞═══════╪═══════╡
│ a     ┆ 1     │
│ a     ┆ 3     │
│ a     ┆ 5     │
│ b     ┆ 10    │
│ b     ┆ 10    │
│ b     ┆ 10    │
│ b     ┆ 20    │
└───────┴───────┘
""")

期望输出仅包含fruit为b的行（该组有4个元素）：

┌───────┬───────┐
│ fruit ┆ price │
│ ---   ┆ ---   │
│ str   ┆ i64   │
╞═══════╪═══════╡
│ b     ┆ 50    │
└───────┴───────┘

解决方案 #

Polars没有直接对应min_count的参数，可通过分组聚合+过滤的两步逻辑实现：

1. 分组后同时计算目标列的求和值，以及组内元素的计数；
2. 过滤出计数≥指定阈值的分组；
3. 移除计数列，得到最终结果。

代码实现 #

result = (
    df
    .group_by("fruit")
    .agg(
        pl.col("price").sum().alias("price"),
        pl.col("price").count().alias("group_count")
    )
    .filter(pl.col("group_count") >= 4)
    .drop("group_count")
)

print(result)

补充：窗口函数实现方式 #

也可以先通过窗口函数标记每组的元素数量，再过滤后聚合，适合需要保留原始行信息的场景，但性能略低于先聚合后过滤的方式：

result = (
    df
    .with_columns(pl.col("fruit").count().over("fruit").alias("group_count"))
    .filter(pl.col("group_count") >= 4)
    .group_by("fruit")
    .agg(pl.col("price").sum().alias("price"))
)

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