如何将pandas的.corrWith两DataFrame相关性计算逻辑转换为PySpark实现

PySpark没有内置和pandas corrwith完全对齐的API，可通过PySpark自带的stat.corr方法手动实现，和原pandas逻辑输出完全一致。

前置说明 #

• 两个DataFrame行顺序必须严格对应，否则会导致相关性计算错误
• 默认计算皮尔逊相关系数，和pandas corrwith默认参数对齐，可修改method参数调整为斯皮尔曼等其他相关系数

实现代码 #

1. 安全版本（带行对齐校验，避免行错位） #

推荐使用该版本，先给两个DataFrame添加统一行号保证行匹配：

from pyspark.sql import functions as F

# 步骤1：给两个DataFrame添加行号，保证按行对齐
x = x.withColumn("row_id", F.monotonically_increasing_id())
y = y.withColumn("row_id", F.monotonically_increasing_id())
# 按行号关联两个DF
joined_df = x.join(y, on="row_id", how="inner")

# 步骤2：遍历所有列对计算相关系数
corr_data = []
x_cols = [col for col in x.columns if col != "row_id"]
y_cols = [col for col in y.columns if col != "row_id"]

for x_col in x_cols:
    for y_col in y_cols:
        # 第三个参数可指定相关系数类型：pearson/spearman
        corr_val = joined_df.stat.corr(x_col, y_col, method="pearson")
        corr_data.append((x_col, y_col, corr_val))

# 步骤3：转成和pandas corrwith输出格式一致的结果
# 格式说明：行是x的列名，列是y的列名，值为对应相关系数
corr_result = spark.createDataFrame(corr_data, schema=["x_col", "y_col", "corr"]) \
    .groupBy("x_col").pivot("y_col").agg(F.first("corr"))

# 如果需要转成pandas DataFrame，直接调用toPandas即可：
# corr_result_pd = corr_result.toPandas().set_index("x_col")

2. 精简版本（适合确认行顺序一致的场景） #

如果已经确认两个DataFrame行顺序完全匹配，可简化实现：

from pyspark.sql import functions as F

corr_data = []
for x_col in x.columns:
    for y_col in y.columns:
        corr_val = x.crossJoin(y.select(y_col)).stat.corr(x_col, y_col)
        corr_data.append((x_col, y_col, corr_val))

corr_result = spark.createDataFrame(corr_data, schema=["x_col", "y_col", "corr"]) \
    .groupBy("x_col").pivot("y_col").agg(F.first("corr"))

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