计算数据框中的所有“-1”和-1以获得某些列的列表

Question

我有一个极坐标数据框，对于某些列，我想计算“-1”（如果是字符）和-1（如果是数字）的数量。我真的想让它成为一个快速查询，所以我对学习正确的 Polars 代码非常感兴趣，这样它运行得非常快。

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]
                    } 
)

所需的数据框如下所示：

df_desired = pl.from_dict({"groups":["B", "C", "A"], "int_1":[0,1,0], "int_2":[0,0,0], "char_1":[1,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())

然后我可能会连接这两个表，但我想在一个查询中完成此操作。

我还想定义两个列表，一个用于对哪些变量进行分组，另一个是我有兴趣查找

-1s

的变量集，这样我就不必对“组”进行硬编码，但我真的不确定我所做的是否是好的代码：

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())

预期结果：

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

编辑：我想我已经成功了。

# attempt:
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())

老实说我不知道这会起作用。我想我对何时可以使用列选择器缺乏了解......

Answer 1

当您测试“相同值”时，另一种方法可能是

.cast()

所有列。

转换为 int：

(df.group_by(group_cols)
   .agg((pl.col(cols_to_check).cast(int, strict=False) == -1).sum())
)

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

转换为字符串：

(df.group_by(group_cols)
   .agg((pl.col(cols_to_check).cast(pl.String) == "-1").sum())
)