如何高效并行对宽DataFrame多列独立运行Logistic Regression？

我有一个包含大量标签列的宽DataFrame（结构如下），需要为每一列独立运行Logistic Regression，正在寻找最高效的并行实现方式。

数据结构 #

+----------+--------+--------+--------+-----+------------+
| features | label1 | label2 | label3 | ... | label30000 |
+----------+--------+--------+--------+-----+------------+

我尝试的方法（性能不佳） #

我最初用ThreadPoolExecutor来并行处理每一列，获取结果后关联，但运行速度很慢：

extract_prob = udf(lambda x: float(x[1]), FloatType())

def lr_for_column(argm):
    col_name = argm[0]
    test_res = argm[1]
    lr = LogisticRegression(featuresCol="features", labelCol=col_name, regParam=0.1)
    lrModel = lr.fit(tfidf)
    res = lrModel.transform(test_tfidf)
    test_res = test_res.join(res.select('id', 'probability'), on="id")
    test_res = test_res.withColumn(col_name, extract_prob('probability')).drop("probability")
    return test_res.select('id', col_name)

with futures.ThreadPoolExecutor(max_workers=100) as executor:
    future_results = [executor.submit(lr_for_column, [colname, test_res]) for colname in list_of_label_columns]
    futures.wait(future_results)
    for future in future_results:
        test_res = test_res.join(future.result(), on="id")

请问有没有更快速的实现方式？

你的问题核心在于没有利用Spark本身的分布式计算能力，反而在Driver端用ThreadPoolExecutor做并行——这会导致所有训练任务都挤在Driver节点，不仅浪费了集群的分布式资源，还可能因为Driver内存/CPU瓶颈拖慢整体速度。下面给你几个更高效的方案：

方案1：将宽表转为长表（Melt），利用Spark分布式分组训练 #

把多标签列转成(id, features, label_name, label_value)的长格式，然后按label_name分组，每组训练一个LR模型，最后再把结果转成宽表。这种方式完全利用Spark的分布式集群资源，每个标签的训练任务会被分配到不同的Executor节点。

实现步骤： #

1. 宽表转长表：用stack函数把所有标签列转成行：

from pyspark.sql.functions import expr, col

# 构造stack表达式：stack(N, 'label1', label1, 'label2', label2, ...)
stack_expr = f"stack({len(list_of_label_columns)}, " + \
             ", ".join([f"'{col}', {col}" for col in list_of_label_columns]) + \
             ") as (label_name, label_value)"

long_df = tfidf.select("id", "features", expr(stack_expr))

2. 分组训练+预测：用groupBy("label_name")结合applyInPandas来训练每个标签的LR模型：

from pyspark.ml.classification import LogisticRegression
import pandas as pd

def train_lr_per_group(pdf):
    label_name = pdf["label_name"].iloc[0]
    # 训练LR模型
    lr = LogisticRegression(featuresCol="features", labelCol="label_value", regParam=0.1)
    model = lr.fit(pdf)
    # 预测概率（取正类概率）
    pred_pdf = model.transform(pdf)
    pred_pdf[label_name] = pred_pdf["probability"].apply(lambda x: x[1])
    return pred_pdf[["id", label_name]]

# 用applyInPandas分组处理，利用Spark分布式计算
result_long = long_df.groupBy("label_name").applyInPandas(train_lr_per_group, schema="id string, label_value double")

# 转回宽表
result_wide = result_long.groupBy("id").pivot("label_name").max("label_value")

方案2：利用Spark的并行模型训练API（Spark 3.0+） #

Spark 3.0及以上对多标签场景有更好的支持，你也可以借助Pipeline结合分布式训练的思路，核心是让训练任务下沉到Executor节点，而不是在Driver端串行/线程池并行。

另外，你之前的代码中每次训练都执行join，会产生大量Shuffle操作，这也是性能瓶颈之一。上面的长表转宽表方式只需要一次最终的pivot，Shuffle次数更少，效率更高。

方案3：优化Driver端并行（如果必须用线程池） #

如果你因为某些限制必须在Driver端并行，那至少要优化这几个点：

• 不要把整个test_resDataFrame传给每个线程，而是让每个线程直接从Spark读取数据，避免序列化/反序列化的巨大开销
• 用spark.sparkContext.parallelize把标签列列表转为RDD，然后用map操作分布式执行训练任务，这比ThreadPoolExecutor更适配Spark场景：

def train_lr(label_col):
    lr = LogisticRegression(featuresCol="features", labelCol=label_col, regParam=0.1)
    model = lr.fit(tfidf)
    pred = model.transform(test_tfidf)
    return pred.select("id", col("probability").apply(lambda x: x[1]).alias(label_col))

# 并行训练每个标签
label_rdd = spark.sparkContext.parallelize(list_of_label_columns, numSlices=100)
result_dfs = label_rdd.map(train_lr).collect()

# 合并所有结果
test_res = test_res.select("id")
for df in result_dfs:
    test_res = test_res.join(df, on="id")

关键性能优化点总结 #

• 避免Driver端瓶颈：尽量把训练任务放到Executor节点执行，充分利用Spark分布式集群的算力
• 减少Shuffle操作：长表转宽表的方式比多次join更高效，能大幅降低IO开销
• 复用数据：提前对tfidf和test_tfidf执行cache()或persist()，避免重复计算特征数据

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