如何用Python-Polars统计符合0-10区间的True实例数量？

我有如下结构的Polars DataFrame：

df = pl.DataFrame({
    "datetime": [
        "2024-09-24 00:00",
        "2024-09-24 01:020",
        "2024-09-24 02:00",
        "2024-09-24 03:00",
    ],
    "Bucket1": [2.5, 8, 0.7, 12],
    "Bucket2": [3.7, 10.1, 25.9, 9.9],
    "Bucket3": [40.0, 15.5, 10.7, 56],
})

我需要统计每行中Bucket1、Bucket2、Bucket3列里值处于0-10左闭区间的数量，预期输出如下：

shape: (4, 2)
┌───────────────────┬──────┐
│ datetime          ┆ 0-10 │
│ ---               ┆ ---  │
│ str               ┆ u32  │
╞═══════════════════╪══════╡
│ 2024-09-24 00:00  ┆ 2    │
│ 2024-09-24 01:020 ┆ 1    │
│ 2024-09-24 02:00  ┆ 1    │
│ 2024-09-24 03:00  ┆ 1    │
└───────────────────┴──────┘

我尝试了两种方法：

• 使用pl.when结合.is_between，但结果仅返回1，无法统计符合条件的Bucket数量；
• 使用concat_list，代码如下：

columns = ["Bucket1", "Bucket2", "Bucket3"]
df.with_columns(
    pl.concat_list(
        [pl.col(col).is_between(0,10,closed="left") for col in columns]
    )
    .arr.sum()
    .alias("0-10")
)

但报错：Invalid input for "col", Expected iterable of type "str" or "DataType", got iterable of "Expr"。请问如何用Polars解决该问题？

方法一：使用pl.sum_horizontal（推荐） #

Polars的sum_horizontal支持对多个表达式进行逐行求和，布尔值会自动转换为1（True）和0（False），直接对每个Bucket列的区间判断结果求和即可：

columns = ["Bucket1", "Bucket2", "Bucket3"]
result = df.with_columns(
    pl.sum_horizontal(
        pl.col(col).is_between(0, 10, closed="left") for col in columns
    ).alias("0-10")
).select("datetime", "0-10")

print(result)

方法二：修复concat_list的用法 #

如果坚持使用concat_list，可以先为每个Bucket列生成布尔判断列，再合并为列表后求和：

columns = ["Bucket1", "Bucket2", "Bucket3"]
result = df.with_columns(
    pl.concat_list(
        [pl.col(col).is_between(0,10,closed="left").alias(f"{col}_bool") for col in columns]
    ).arr.sum().alias("0-10")
).select("datetime", "0-10")

print(result)

方法三：使用pl.row逐行处理 #

通过pl.row遍历每行的指定列，统计符合条件的数值数量：

columns = ["Bucket1", "Bucket2", "Bucket3"]
result = df.with_columns(
    pl.row(columns).map_elements(
        lambda row: sum(0 <= x < 10 for x in row),
        return_dtype=pl.UInt32
    ).alias("0-10")
).select("datetime", "0-10")

print(result)

三种方法都能得到预期输出，其中sum_horizontal是性能最优的方案，它利用Polars的向量化操作，避免了逐行遍历的额外开销。

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