pyspark 中的重型有状态 UDF

Posted 2023-04-13

【中文标题】pyspark 中的重型有状态 UDF

【英文标题】：Heavy stateful UDF in pyspark

【发布时间】：2018-12-03 10:54:41

【问题描述】：

我必须在 Spark 中作为 UDF 运行一个非常繁重的 python 函数，并且我想在 UDF 中缓存一些数据。案例类似here@

我知道这是缓慢且错误的。 但是现有的基础设施已经启动，我不想建立一个新的基础设施并单独处理这种情况下的数据加载/并行化/故障安全。

这就是我的 spark 程序的样子：

from mymodule import my_function # here is my function
from pyspark.sql.types import *
from pyspark.sql.functions import udf
from pyspark.sql.session import SparkSession

spark = SparkSession.builder.getOrCreate()

schema = StructType().add("input", "string")
df = spark.read.format("json").schema(schema).load("s3://input_path")

udf1 = udf(my_function, StructType().add("output", "string"))
df.withColumn("result", udf1(df.input)).write.json("s3://output_path/")

my_function 在内部调用具有慢构造函数的对象的方法。 因此我不希望为每个条目初始化对象，我正在尝试缓存它：

from my_slow_class import SlowClass
from cachetools import cached

@cached(cache=) 
def get_cached_object():
    # this call is really slow therefore I am trying 
    # to cache it with cachetools
    return SlowClass() 

def my_function(input):
    slow_object = get_cached_object()
    output = slow_object.call(input)
    return 'output': output

mymodule 和 my_slow_class 作为模块安装在每台 spark 机器上。

它似乎工作。构造函数只被调用了几次（输入数据帧中的 100k 行只调用了 10-20 次）。这就是我想要的。

我担心的是 Spark 执行器内部的多线程/多处理，以及缓存的 SlowObject 实例是否在许多并行 my_function 调用之间共享。

我可以相信my_function 在工作节点上的python 进程中一次调用一次这一事实吗？ spark是否在执行我的UDF的python进程中使用任何多处理/多线程？

【参考方案1】：

Spark 分叉 Python 进程以创建单个工作进程，但单个工作进程中的所有处理都是顺序的，除非UserDefinedFunction 明确使用多线程或多处理。

所以只要状态用于缓存并且slow_object.call 是一个纯函数，你就不用担心了。