在PySpark UDF中记录Worker节点ID或唯一标识的实现方法

方法1：获取Worker节点IP地址（最通用） #

直接在UDF内调用Python标准库socket即可获取当前执行节点的IP/主机名，不需要额外依赖，示例代码如下：

import socket
from pyspark.sql.types import StringType
from pyspark.sql.functions import udf

@udf(returnType=StringType())
def your_compute_udf(your_params):
    # 记录节点标识
    worker_hostname = socket.gethostname()
    worker_ip = socket.gethostbyname(worker_hostname)
    # 原有计算、统计逻辑保持不变
    # ...
    return your_result

说明：主机名或IP都是Worker维度的唯一标识，完全满足任务分配排查的需求

方法2：获取Spark原生Worker/Executor ID #

Spark启动Executor进程时会将节点信息写入环境变量，你可以直接读取获取官方分配的ID：

import os
from pyspark.sql.types import StringType
from pyspark.sql.functions import udf

@udf(returnType=StringType())
def your_compute_udf(your_params):
    # 读取Spark环境变量，设置默认值避免特殊场景报错
    worker_id = os.environ.get("SPARK_WORKER_ID", "unknown_worker")
    executor_id = os.environ.get("SPARK_EXECUTOR_ID", "unknown_executor")
    # 原有计算、统计逻辑保持不变
    # ...
    return your_result

说明：Executor ID是进程维度的标识，同一个Worker上可能跑多个Executor，结合Worker ID可以做更细粒度的调度排查

额外排查建议 #

你当前仅处理100行数据，出现部分Worker未被占用的情况大概率是分区数太少：Spark默认会按数据量生成分区，100行通常只会生成1-2个分区，每个分区只会被分配到一个Executor上运行，自然不会调度到所有Worker。如果要充分利用所有节点算力，可以先调用df = df.repartition(2 * 集群总CPU核心数)将数据打散成足够多的分区，再运行UDF即可。

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