Apache Spark中如何按文件路径将DataFrame任务分配到指定节点？

集群配置：两台PC通过以太网连接，PC A同时作为Master和Worker节点，PC B仅为Worker节点。因条件限制无法使用分布式存储系统，需从本地目录/root/data创建DataFrame——该目录在A、B节点均存在，但A节点下是1.txt至2000.txt，B节点下是2001.txt至4000.txt。

任务目标：用Pandas UDF统计每个文件的字符数，生成4000行的DataFrame，要求1-2000.txt在A节点处理，2001-4000.txt在B节点处理。

当前问题：现有代码运行时任务会在A、B节点混排，导致节点处理不存在的文件报错，需基于单个DataFrame提交集群任务，并按路径定向分配任务到对应节点（暂不考虑容错）。

现有代码：

spark = SparkSession.builder.config(conf=conf).appName("wordcount").getOrCreate()

file_paths = [f"/root/data/{i}.txt" for i in range(1, 4001)]

data = [(path,) for path in file_paths]
df = spark.createDataFrame(data, ["path"])

df = df.withColumn("count", my_udf_wordCount(df['path']))

save_path = "/root/data/result"
df.write.format('com.databricks.spark.csv') \
        .mode('overwrite') \
        .option("header", "true") \
        .save(save_path)

@pandas_udf(IntegerType())
def my_udf_wordCount(iterator: Iterator[pd.Series]) -> Iterator[pd.Series]:
    for path_series in iterator:
        for path in path_series:
            file_path = path.strip()
            with open(file_path, 'r') as file:
                text = file.read()
                words = text.split()
                word_count = len(words)
                yield pd.Series(word_count)

要实现任务定向分配，核心是利用Spark的节点标签和任务本地属性，结合数据分区来控制任务运行节点：

1. 给集群节点添加标签 #

先给两个节点设置唯一标签，方便Spark识别：

• PC A（Master+Worker）：添加标签 node=nodeA
• PC B（Worker）：添加标签 node=nodeB

如果是Standalone集群，可通过Web UI（默认http://master:8080）的节点管理页面添加标签；也可修改spark-defaults.conf或启动Worker时通过参数指定。

2. 给DataFrame添加节点标识列 #

从文件路径中提取文件编号，判断该文件归属的节点，新增node_label列：

from pyspark.sql.functions import when, col

# 从路径中提取文件数字（比如"/root/data/1234.txt"提取1234）
df = df.withColumn("file_num", col("path").substr(-8, 4).cast("int"))
# 根据编号判断归属节点
df = df.withColumn("node_label", 
                   when(col("file_num") <= 2000, "nodeA")
                   .otherwise("nodeB"))

3. 分区处理并指定运行节点 #

将DataFrame按节点标识拆分，分别设置任务的运行节点标签，确保任务只分配到对应节点：

# 拆分数据为A、B节点对应的子集
node_a_df = df.filter(col("node_label") == "nodeA")
node_b_df = df.filter(col("node_label") == "nodeB")

# 定义带节点标签的处理函数
def process_on_node(df, target_node):
    # 设置当前任务必须运行在指定标签的节点上
    spark.sparkContext.setLocalProperty("spark.scheduler.nodeLabel", target_node)
    # 应用UDF处理数据
    return df.withColumn("count", my_udf_wordCount(df['path']))

# 分别处理两个子集
processed_a = process_on_node(node_a_df, "nodeA")
processed_b = process_on_node(node_b_df, "nodeB")

# 合并处理结果
final_df = processed_a.union(processed_b)

4. 保存结果 #

保持原有保存逻辑即可：

save_path = "/root/data/result"
final_df.write.format('com.databricks.spark.csv') \
        .mode('overwrite') \
        .option("header", "true") \
        .save(save_path)

关键注意事项 #

• 确保节点标签配置正确，Spark集群能识别标签与节点的映射；
• Standalone集群需在spark-defaults.conf中启用节点标签功能：spark.scheduler.nodeLabelExpression.enabled true；
• 通过setLocalProperty设置的节点标签仅对当前线程的任务生效，拆分处理能保证不同子集的任务定向到对应节点。

内容的提问来源于stack exchange，提问作者이하운