Pandas按年份统计含特定Code的文档数问题及Polars实现咨询

问题回顾 #

现有DataFrame结构如下：

year docnr   code
2000 doc1 T1
2000 doc1 T2
2000 doc2 F1
2000 doc2 F1
2000 doc2 F1
2001 doc3 T1
2001 doc2 F1
2001 doc2 F1

需求：按year和docnr分组，统计每年中至少包含一个含"T"的code的文档数量，最终输出以年份为索引、统计数为列的DataFrame（预期2000年为1，2001年为1）。

Pandas解决方案 #

你之前的代码问题在于any("T" in x)的判断逻辑错误：x是分组后的Series，"T" in x是检查"T"是否属于Series的索引，而非元素中包含"T"。正确的做法是先判断每个code是否包含"T"，再检查该文档是否存在符合条件的记录，最后按年份统计数量。

完整代码：

import pandas as pd

# 构造示例数据
data = {
    "year": [2000, 2000, 2000, 2000, 2000, 2001, 2001, 2001],
    "docnr": ["doc1", "doc1", "doc2", "doc2", "doc2", "doc3", "doc2", "doc2"],
    "code": ["T1", "T2", "F1", "F1", "F1", "T1", "F1", "F1"]
}
df = pd.DataFrame(data)

# 第一步：按year+docnr分组，标记每个文档是否含"T"的code
doc_has_T = df.groupby(["year", "docnr"])["code"].apply(lambda x: x.str.contains("T").any()).reset_index(name="has_T")

# 第二步：按year分组，统计含"T"的文档数量
result = doc_has_T.groupby("year")["has_T"].sum().to_frame(name="count")

print(result)

输出结果：

count
year       
2000      1
2001      1

Polars解决方案 #

Polars的语法更简洁，无需多步分组，通过链式调用即可完成，且原生支持向量化操作，处理大数据集时性能更优：

import polars as pl

# 构造示例数据
data = {
    "year": [2000, 2000, 2000, 2000, 2000, 2001, 2001, 2001],
    "docnr": ["doc1", "doc1", "doc2", "doc2", "doc2", "doc3", "doc2", "doc2"],
    "code": ["T1", "T2", "F1", "F1", "F1", "T1", "F1", "F1"]
}
df = pl.DataFrame(data)

# 链式调用完成统计
result = (
    df
    .group_by(["year", "docnr"])
    .agg(pl.col("code").str.contains("T").any().alias("has_T"))
    .group_by("year")
    .agg(pl.col("has_T").sum().alias("count"))
    .set_index("year")
)

print(result)

输出结果：

shape: (2, 1)
┌──────┬───────┐
│ year ┆ count │
│ ---  ┆ ---   │
│ i64  ┆ u32   │
╞══════╪═══════╡
│ 2000 ┆ 1     │
│ 2001 ┆ 1     │
└──────┴───────┘

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