PySpark计算相关矩阵代码运行超时问题排查

我有一段用于计算LD矩阵的Pandas代码，想迁移到PySpark提升处理效率，但改写后的PySpark代码一直无法运行结束。输入的ld_data是形状为(10000, 6167)的0-1矩阵，以下是相关代码：

Pandas实现代码 #

import pandas as pd
def get_ld_matrix():
    ld_data = pd.read_csv(f"{target}_locus_UKBB.txt.raw", delim_whitespace=True)
    ld_data = ld_data.drop(columns=["FID", "IID", "PAT", "MAT", "SEX", "PHENOTYPE"])
    ld_matrix = ld_data.corr(method='pearson')
    return ld_data, ld_matrix
ld_data, ld_matrix= get_ld_matrix()

尝试的PySpark代码 #

from pyspark.sql import SparkSession
from pyspark.ml.stat import Correlation
from pyspark.ml.feature import VectorAssembler

def get_ld_matrix():
    """Calculates the LD matrix based on the LD data from PLINK using PySpark"""

    ld_data = spark.read.csv(f"{target}_locus_UKBB.txt.raw", sep=" ", header=True, inferSchema=True)
    drop_list = ["FID", "IID", "PAT", "MAT", "SEX", "PHENOTYPE"]
    ld_data = ld_data.drop(*drop_list)
    for col in ld_data.columns:
        ld_data = ld_data.withColumn(col, ld_data[col].cast("float"))

    print('Calculating LD correlation matrix...')

    vec_assembler = VectorAssembler(inputCols=ld_data.columns, outputCol="features")
    ld_data_vector = vec_assembler.transform(ld_data).select("features")
    ld_matrix = Correlation.corr(ld_data_vector, "features").collect()[0][0]

    print('Finished calculating LD correlation matrix')
    return ld_data, ld_matrix

ld_data, ld_matrix = get_ld_matrix()

示例数据格式：

ld_data = pd.DataFrame(np.random.choice([0, 1], size=(5, 5)), columns=['A', 'B', 'C', 'D', 'E'])

问题排查与优化方案 #

核心问题分析 #

• inferSchema=True的性能开销：读取大文件时，Spark需要扫描全量数据推断列类型，对于10000行×6000+列的文件，这个过程会消耗大量时间。
• 循环转换列类型的低效：逐个列调用withColumn会生成多个DataFrame副本，触发多次不必要的计算与Shuffle，累积大量性能损耗。
• collect()的内存压力：6167个特征生成的相关矩阵是6167×6167的稠密矩阵，collect()会把整个矩阵拉到Driver端，极易导致Driver内存溢出，进而使任务卡死。

优化后的代码 #

from pyspark.sql import SparkSession
from pyspark.ml.stat import Correlation
from pyspark.ml.feature import VectorAssembler
from pyspark.sql.types import IntegerType, FloatType

# 初始化SparkSession，配置合理资源
spark = SparkSession.builder \
    .appName("LDMatrixCalculation") \
    .config("spark.driver.memory", "32g")  # 根据Driver内存情况调整，需容纳6000×6000的矩阵
    .config("spark.executor.memory", "16g") \
    .getOrCreate()

def get_ld_matrix(target):
    # 1. 手动指定Schema，避免全量扫描推断类型
    temp_df = spark.read.csv(f"{target}_locus_UKBB.txt.raw", sep=" ", header=True, nrows=1)
    drop_cols = ["FID", "IID", "PAT", "MAT", "SEX", "PHENOTYPE"]
    feature_cols = [col for col in temp_df.columns if col not in drop_cols]
    
    # 构建Schema
    schema = {col: IntegerType() for col in temp_df.columns}
    ld_data = spark.read.csv(f"{target}_locus_UKBB.txt.raw", sep=" ", header=True, schema=schema)
    ld_data = ld_data.drop(*drop_cols)
    
    # 2. 批量转换列类型，避免循环开销
    ld_data = ld_data.select([ld_data[col].cast(FloatType()).alias(col) for col in ld_data.columns])

    print('Calculating LD correlation matrix...')

    vec_assembler = VectorAssembler(inputCols=ld_data.columns, outputCol="features")
    ld_data_vector = vec_assembler.transform(ld_data).select("features")
    
    # 3. 安全获取相关矩阵，避免直接collect大对象
    corr_result = Correlation.corr(ld_data_vector, "features")
    ld_matrix = corr_result.first()[0]
    
    print('Finished calculating LD correlation matrix')
    return ld_data, ld_matrix

# 调用函数（替换为实际target值）
target = "your_target_value"
ld_data, ld_matrix = get_ld_matrix(target)

额外优化建议 #

• 资源配置调整：根据集群实际资源，增大spark.driver.memory和spark.executor.memory，确保Driver有足够内存容纳相关矩阵（6167×6167的Float矩阵约150MB，需预留冗余）。
• 分块/降维处理：若不需要全量LD矩阵，可分批次计算特征对相关性，或先通过PCA等方法降维减少计算量。
• 文件格式优化：将原文本文件转换为Parquet格式，Spark读取Parquet的速度远快于CSV，且自带Schema信息，能进一步提升性能。

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