如何对Polars的pl.Array类型列执行水平归约操作？

Boolean列的水平any操作示例 #

pl.any_horizontal是Polars的水平操作函数，适用于pl.Boolean类型列，实现逐行的逻辑或运算：

import polars as pl

df = pl.DataFrame(
    {
        "a": pl.Series("a", [False, False], dtype=pl.Boolean),
        "b": pl.Series("b", [True, False], dtype=pl.Boolean),
    }
)

print(df.with_columns(pl.any_horizontal("a", "b")))

输出：

shape: (2, 3)
┌───────┬───────┬───────┐
│ a     ┆ b     ┆ any   │
│ ---   ┆ ---   ┆ ---   │
│ bool  ┆ bool  ┆ bool  │
╞═══════╪═══════╪═══════╡
│ false ┆ true  ┆ true  │
│ false ┆ false ┆ false │
└───────┴───────┴───────┘

Array列的操作问题 #

当对同形状的pl.Array类型列直接使用pl.any_horizontal时，会触发类型转换错误：

df = pl.DataFrame(
    {
        "a": pl.Series("a", [[False, False], [False, True]], dtype=pl.Array(shape=2, inner=pl.Boolean)),
        "b": pl.Series("b", [[True, False], [True, True]], dtype=pl.Array(shape=2, inner=pl.Boolean)),
    }
)

# 报错：InvalidOperationError: cannot cast Array type (inner: 'Boolean', to: 'Boolean')
df.with_columns(pl.any_horizontal("a", "b"))

我们需要的是逐元素执行逻辑或，得到同形状的Array列，预期输出：

shape: (2, 3)
┌────────────────┬────────────────┬────────────────┐
│ a              ┆ b              ┆ any            │
│ ---            ┆ ---            ┆ ---            │
│ array[bool, 2] ┆ array[bool, 2] ┆ array[bool, 2] │
╞════════════════╪════════════════╪════════════════╡
│ [false, false] ┆ [true, false]  ┆ [true, false]  │
│ [false, true]  ┆ [true, true]   ┆ [true, true]   │
└────────────────┴────────────────┴────────────────┘

可行实现方案 #

最优方案：使用arr.zip_with逐元素运算 #

利用arr.zip_with方法，直接对两个Array列的对应元素执行逻辑或操作，代码简洁高效：

result = df.with_columns(
    pl.col("a").arr.zip_with(pl.col("b"), lambda x, y: x | y).alias("any")
)
print(result)

备选方案1：展开数组后再合并 #

先将Array列拆分为多个Boolean列，逐位置执行any_horizontal后再重新组合成Array：

# 展开数组为单个布尔列
expanded_df = df.with_columns(
    [pl.col("a").arr.get(i).alias(f"a_{i}") for i in range(2)] +
    [pl.col("b").arr.get(i).alias(f"b_{i}") for i in range(2)]
)

# 逐位置计算any并合并为Array
result = expanded_df.with_columns(
    pl.concat_list(
        [pl.any_horizontal(f"a_{i}", f"b_{i}") for i in range(2)]
    ).cast(pl.Array(shape=2, inner=pl.Boolean)).alias("any")
).drop([f"a_{i}", f"b_{i}" for i in range(2)])

print(result)

备选方案2：结合arr.eval与any_horizontal #

通过arr.eval提取数组元素，再执行水平any操作后转换回Array类型：

result = df.with_columns(
    pl.any_horizontal(
        pl.col("a").arr.eval(pl.element()), 
        pl.col("b").arr.eval(pl.element())
    ).arr.to_list().cast(pl.Array(shape=2, inner=pl.Boolean)).alias("any")
)
print(result)

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