如何使用Polars检测问卷DataFrame中缺失的答案选项并优化实现方案

嘿，你的代码已经能完美完成需求啦！不过咱们可以用更符合Polars idiom的方式来实现——也就是尽量利用Polars的向量化操作和原生API，避免手动循环列，让代码更简洁、高效，也更易维护。

优化后的代码 #

import polars as pl

# 定义预期的答案选项（和你原来的一致）
fields = {
    'q1': [f'Q1A{i}' for i in range(1,12)],
    'q2': [f'Q2A{i}' for i in range(1,3)],
    'q3': [f'Q3A{i}' for i in range(1,8)],
    'q4': [f'Q4A{i}' for i in range(1,4)]
}

# 原始问卷数据
df = pl.DataFrame({
    'id': [1,2,3,4,5],
    'q1': ['Q1A1,Q1A2','Q1A6','Q1A7,Q1A9','Q1A11','Q1A5,Q1A3'],
    'q2': ['Q2A1','','Q2A2','Q2A1','Q2A1'],
    'q3': ['Q3A1','Q3A3','Q3A2,Q3A5','Q3A6','Q3A7'],
    'q4': ['Q4A1,Q4A3','Q4A1','','Q4A1','Q4A1'],
})

# 步骤1：把预期选项转成结构化的长表
expected_options = pl.concat([
    pl.DataFrame({
        "question": question,
        "option": options
    }) for question, options in fields.items()
])

# 步骤2：处理原始数据，提取实际出现的非空选项
actual_options = df.melt(id_vars=["id"], value_name="raw_options") \
    .filter(pl.col("variable").str.starts_with("q")) \
    .with_columns(pl.col("raw_options").str.split(",")) \
    .explode("raw_options") \
    .filter(pl.col("raw_options") != "") \
    .select(pl.col("variable").alias("question"), pl.col("raw_options").alias("option")) \
    .unique()

# 步骤3：用anti join找出预期存在但实际缺失的选项
missing_options = expected_options.join(
    actual_options,
    on=["question", "option"],
    how="anti"
)

# 步骤4：聚合结果，补充无缺失的问题
final_result = missing_options.group_by("question") \
    .agg(
        pl.count().alias("missing_count"),
        pl.col("option").str.join(", ").alias("missing_options")
    ) \
    .join(
        pl.DataFrame({"question": list(fields.keys())}),
        on="question",
        how="outer"
    ) \
    .with_columns(
        pl.col("missing_count").fill_null(0),
        pl.col("missing_options").fill_null("")
    ) \
    .sort("question")

# 输出结果（和你原代码的格式匹配）
for row in final_result.rows():
    print(row)

运行结果 #

('q1', 3, 'Q1A4, Q1A8, Q1A10')
('q2', 0, '')
('q3', 1, 'Q3A4')
('q4', 1, 'Q4A2')

为什么这更“Polars-ish”？ #

1. 避免Python循环：用pl.concat和melt批量处理所有问卷列，充分利用Polars的向量化计算优势，比手动循环列效率更高（尤其是数据量大的时候）。
2. 链式调用更清晰：每个步骤都用Polars原生API串联，逻辑一目了然——从转预期表、处理实际数据，到找缺失、聚合结果，整个流程连贯易读。
3. 自动补充无缺失的问题：通过outer join自动把没有缺失选项的问题也纳入结果，不用手动处理空值情况。
4. 可扩展性更强：如果后续新增问卷问题，只需要在fields里添加对应选项，代码逻辑完全不用改，维护成本更低。

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