如何在Polars的df.corr方法中使用自定义函数（类Pandas方式）

Polars的df.corr()方法目前仅支持内置的Pearson、Spearman等统计方法，不支持直接传入自定义函数。但我们可以通过构造列对的批量计算逻辑，实现类似Pandas中df.corr(method=自定义函数)的功能，同时保留Polars在大数据集上的性能优势。

方案1：构造列对批量计算 #

核心思路是生成所有需要计算的列组合，通过分组计算自定义统计量后，整理成相关矩阵格式：

from sklearn.datasets import load_iris
import polars as pl
from scipy.stats import pearsonr

def pearsonr_pval(x: pl.Series, y: pl.Series) -> float:
    return pearsonr(x.to_numpy(), y.to_numpy())[1]

# 加载数据集
df = pl.DataFrame(load_iris()['data'], schema=load_iris()['feature_names'])
cols = df.columns

# 生成所有列的两两组合
column_pairs = [(col1, col2) for col1 in cols for col2 in cols]

# 构造列对映射表并交叉连接原数据
pair_df = pl.DataFrame({"col1": [p[0] for p in column_pairs], "col2": [p[1] for p in column_pairs]})
cross_df = pair_df.join(df, how="cross")

# 定义分组计算函数
def compute_pval(group: pl.DataFrame) -> pl.DataFrame:
    col1, col2 = group["col1"][0], group["col2"][0]
    p_val = pearsonr_pval(group[col1], group[col2])
    return pl.DataFrame({"col1": [col1], "col2": [col2], "p_val": [p_val]})

# 分组计算并转成矩阵格式
p_val_results = cross_df.group_by(["col1", "col2"]).map_groups(compute_pval)
p_val_matrix = p_val_results.pivot(index="col1", columns="col2", values="p_val")

print(p_val_matrix)

方案2：轻量循环+向量化操作 #

针对大数据集，交叉连接可能占用较多内存，改用列对循环结合Polars的快速数组转换，减少内存开销：

from sklearn.datasets import load_iris
import polars as pl
from scipy.stats import pearsonr

df = pl.DataFrame(load_iris()['data'], schema=load_iris()['feature_names'])
cols = df.columns

# 初始化结果存储字典
p_val_dict = {col: [] for col in cols}

# 遍历列对计算p值
for col1 in cols:
    col1_data = df[col1].to_numpy()
    for col2 in cols:
        p_val = pearsonr(col1_data, df[col2].to_numpy())[1]
        p_val_dict[col2].append(p_val)

# 转换为Polars矩阵格式
p_val_matrix = pl.DataFrame(p_val_dict, index=cols)

print(p_val_matrix)

性能提示 #

虽然自定义函数无法完全利用Polars内置算法的极致优化，但Polars的to_numpy()操作比Pandas更高效，且内存管理机制更优，在大数据集上仍能获得比Pandas更快的处理速度。若要进一步优化，可尝试将自定义函数改为直接支持Polars Series的向量化实现，或使用Polars的map_batches接口批量处理。

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