如何在Databricks上并行执行多参数函数以充分利用Spark集群

为什么joblib没法充分利用Spark集群？ #

joblib本质是单节点内的多核并行，只能用到Driver或者单个Executor节点的CPU资源，没法把任务分发到整个Spark集群的多节点Executor上，自然发挥不了集群的分布式优势。

核心思路：把任务分布式分发到Spark Executor节点 #

利用Spark的RDD/DataFrame分布式计算能力，将每个类别对应的训练任务拆分成独立的子任务，分发到集群的各个Executor节点并行执行。

具体实现步骤 #

1. 准备待处理的类别列表

先从原始数据中提取所有需要单独处理的类别，转成Spark可分布式处理的RDD：

# 从原始Spark DataFrame中获取去重后的类别列表
categories = [row["category"] for row in original_df.select("category").distinct().collect()]
# 转成Spark RDD，每个元素对应一个待处理的类别
categories_rdd = spark.sparkContext.parallelize(categories)

2. 广播共享数据集（优化性能）

如果所有类别任务都依赖同一份原始数据集，用Spark的广播变量把数据集分发到所有Executor节点，避免重复拉取数据：

# 广播原始Spark DataFrame到所有Executor
broadcast_original_df = spark.sparkContext.broadcast(original_df)

3. 编写可分布式执行的训练函数

调整你的训练函数，让它能接收类别参数和广播的数据集，在Executor节点上完成过滤、训练、预测和结果保存：

def train_kmeans(category):
    # 从广播变量获取原始数据，过滤当前类别
    filtered_df = broadcast_original_df.value.filter(f"category = '{category}'")
    # 转换为Pandas DataFrame适配sklearn
    pd_data = filtered_df.toPandas()
    
    # 提取特征（替换成你的实际特征列）
    features = pd_data[["feature_col1", "feature_col2"]]
    
    # 训练sklearn K-Means模型
    from sklearn.cluster import KMeans
    kmeans_model = KMeans(n_clusters=5, random_state=42)
    pd_data["cluster_label"] = kmeans_model.fit_predict(features)
    
    # 转换回Spark DataFrame并保存到表（支持追加模式）
    result_spark_df = spark.createDataFrame(pd_data)
    result_spark_df.write.mode("append").saveAsTable("your_result_table")

4. 触发分布式并行执行

通过RDD的foreach方法，把每个类别任务分发到Executor节点并行运行：

categories_rdd.foreach(train_kmeans)

关键注意事项 #

• Executor环境依赖：确保Spark集群的所有Executor节点都安装了scikit-learn，可以通过Databricks集群的「库」管理界面安装，或者用集群初始化脚本批量安装。
• 资源配置：根据类别数量调整集群的Executor数量和核心数，比如10个类别可以配置5个Executor，每个Executor分配2核，让任务并行度匹配集群资源。
• 数据序列化：广播变量默认用Pickle序列化，确保你的数据集能正常序列化；如果数据量极大，建议先对原始数据做分区优化，减少单节点负载。
• 结果一致性：如果所有任务都写入同一张表，要保证每个任务输出的DataFrame Schema完全一致，避免写入失败。

替代方案：Databricks批量作业调度 #

如果每个类别对应的训练任务计算量极大，可以将每个类别封装成独立的Databricks作业，通过Databricks API批量提交作业，利用多作业集群实现更高程度的并行。这种方式适合任务完全独立、需要单独监控的场景。

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