如何高效统计Polars DataFrame指定列中数值/字符型-1的分组数量

问题描述 #

我有一个Polars DataFrame，需要针对指定列统计数值型-1和字符型"-1"的出现数量，要求实现高效查询逻辑，同时避免硬编码列名。

示例DataFrame #

import polars as pl
import polars.selectors as cs
import numpy as np

df = pl.DataFrame( {"int_1": [1, 2, 3, -1, 5], 
                    "char_1": ["foo", "ham", "spam", "egg", "-1"], 
                    "not_of_interest": np.random.rand(5),
                    "groups": ["A", "A", "B", "C", "B"],
                    "int_2": [12, 12, 13, 14, 15]
                    } 
)

期望结果 #

shape: (3, 4)
┌────────┬───────┬───────┬────────┐
│ groups ┆ int_1 ┆ int_2 ┆ char_1 │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ i64 ┆ i64 │
╞════════╪═══════╪═══════╪════════╡
│ B ┆ 0 ┆ 0 ┆ 1 │
│ C ┆ 1 ┆ 0 ┆ 0 │
│ A ┆ 0 ┆ 0 ┆ 0 │
└────────┴───────┴───────┴────────┘

实现过程与最终方案 #

初始尝试 #

最初我分成两次查询分别处理数值列和字符串列，后续再合并结果：

df.select(cs.numeric() | cs.by_name("groups")) \
.group_by("groups").agg((pl.col("*") == pl.lit(-1)).sum())

df.select(cs.string() | cs.by_name("groups")) \
.group_by("groups").agg((pl.col("*") == pl.lit("-1")).sum())

优化列名管理 #

为避免硬编码列名，我定义了列名列表，但仍需分两次查询：

cols_to_check = ['int_1', 'int_2', 'char_1']
group_cols = ['groups']

df.select((cs.string() & cs.by_name(*cols_to_check)) | cs.by_name(*group_cols)) \
.group_by(group_cols).agg((pl.col("*") == pl.lit("-1")).sum())

df.select((cs.numeric() & cs.by_name(*cols_to_check)) | cs.by_name(*group_cols)) \
.group_by(group_cols).agg((pl.col("*") == pl.lit(-1)).sum())

最终单查询高效方案 #

通过在agg中直接组合列选择器，实现了单条查询完成所有统计：

cols_to_check = ['int_1', 'int_2', 'char_1']
group_cols = ['groups']

df.group_by(group_cols).agg(
    ((cs.numeric() & cs.by_name(*cols_to_check)) == -1).sum(), 
    ((cs.string() & cs.by_name(*cols_to_check)) == "-1").sum()
)

Polars列选择器适用场景 #

• 批量筛选列：一键选中同类型列，比如cs.numeric()选中所有数值列、cs.string()选中所有字符串列，无需手动枚举列名，适配多列场景。
• 组合筛选逻辑：通过&（且）、|（或）、~（非）组合条件，精准定位目标列，比如cs.numeric() & cs.by_name(*cols_to_check)可选中指定列表内的数值列，灵活性远超硬编码。
• 链式操作复用：在select、agg、with_columns等操作中直接嵌入选择器，避免重复定义列列表，让代码更简洁易维护。
• 动态列处理：当列名不确定或需按数据类型动态调整逻辑时，选择器可自动适配，比如针对不同类型列执行差异化统计规则（如本次的数值型与字符串型-1统计）。

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