python - pyspark 中的重型有状态 UDF

Question

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

我知道它很慢而且是错误的。 但现有的基础设施是在 Spark 中，我不想为此情况建立新的基础设施并单独处理数据加载/并行化/故障安全。

这就是我的 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 进程中使用任何多处理/多线程？

Answer 1

Spark fork Python 进程来创建单个工作进程，但是单个工作进程中的所有处理都是连续的，除非 UserDefinedFunction 显式使用多线程或多处理。

因此，只要状态用于缓存并且 slow_object.call 是一个纯函数，您就无需担心。