如何为Polars DataFrame所有数值列高效应用Kruskal-Wallis检验

直接上高效实现，核心是用Polars的melt+分组聚合替代逐列遍历，同时明确UDF的输出类型解决类型推断报错：

实现步骤与代码 #

依赖准备 #

确保安装polars和scipy：

pip install polars scipy

完整代码 #

import polars as pl
from scipy.stats import kruskal
from typing import Tuple

# 定义检验函数，输入分组后的结构化数据，返回统计量和p值
def run_kruskal(group_struct: pl.Series) -> Tuple[float, float]:
    # 拆分struct为Group和value列
    group_data = group_struct.struct.unnest()
    # 按Group提取各组数值（自动过滤空值）
    group_values = [
        group_data.filter(pl.col("Group") == g)["value"].drop_nulls().to_list()
        for g in group_data["Group"].unique().to_list()
    ]
    # 执行检验
    stat, p_val = kruskal(*group_values)
    return (stat, p_val)

# 示例DataFrame
df = pl.DataFrame({
    "Group": ["A", "A", "B", "B", "C", "C"],
    "col1": [1, 2, 3, 4, 5, 6],
    "col2": [10, 20, 30, 40, 50, 60]
})

# 1. 筛选所有数值列（排除Group列）
numeric_cols = [col for col in df.columns if df[col].dtype in pl.NUMERIC_DTYPES and col != "Group"]

# 2. 宽表转长表，将所有数值列合并为一列，保留原列名
melted = df.melt(id_vars="Group", value_vars=numeric_cols, variable_name="column", value_name="value")

# 3. 按原列名分组，将每组的(Group, value)打包为struct
grouped = melted.group_by("column").agg(pl.struct(["Group", "value"]).alias("group_data"))

# 4. 应用检验函数，明确指定返回类型解决报错
results = grouped.with_columns(
    pl.struct(["group_data"]).map_elements(
        lambda x: run_kruskal(x["group_data"]),
        return_dtype=pl.Struct([("kruskal_stat", pl.Float64), ("p_value", pl.Float64)])
    ).alias("test_results")
).unnest("test_results")

print(results)

关键细节 #

• 无遍历批量处理：通过melt把多列数值转为长格式，一次分组即可处理所有列，完全避免逐列循环的低效问题。
• 解决类型推断报错：在map_elements中显式指定return_dtype为结构化类型，Polars无需自动推断，彻底解决「PanicException: UDF failed: BindingsError」问题。
• 空值处理：检验前对每组数据执行drop_nulls()，避免空值导致统计检验失败。
• 性能优势：Polars的分组、聚合操作基于Rust实现，比Pandas循环或转Pandas处理的效率高得多，适合大行数、多列的场景。

输出结果 #

shape: (2, 3)
column  kruskal_stat  p_value
str     f64           f64
---     ---           ---
col1    4.4           0.111111
col2    4.4           0.111111

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