PySpark中对Collect生成的数组列执行线性回归求斜率

问题背景 #

有一个PySpark DataFrame，包含两列Array[Double]类型数据（通过F.collect提取），每行两个数组长度一致。需要以array_1为X（归一化时间戳）、array_2为Y（特征值），对每行执行线性回归计算斜率，新增Slope列并保留原行数。

运行自行编写的代码时触发序列化错误：

Could not serialize object: TypeError: can't pickle _thread.RLock objects

尝试的代码如下：

from pyspark.sql.functions import udf
from pyspark.sql.types import DoubleType
from pyspark.ml.regression import LinearRegression
from pyspark.ml.feature import VectorAssembler

data = [(1, [1, 2, 3, 4], [8, 9, 7, 3]), (2, [5, 6, 7, 8], [1, 2, 3, 4]), (3, [9, 10, 11, 12], [1, 2, 3, 4])]
df = spark.createDataFrame(data, ["id", "array_1", "array_2"])

def calculate_slope(arr1, arr2):
    vector_assembler = VectorAssembler(inputCols=["arr1"], outputCol="features")
    lr = LinearRegression(featuresCol="features", labelCol="arr2")

    def calculate_slope_internal(row):
        temp_data = [(row.id, row.array1, row.array2)]
        temp_df = spark.createDataFrame(temp_data, ["id", "arr1", "arr2"])

        temp_df = vector_assembler.transform(temp_df)
        model = lr.fit(temp_df)

        return float(model.coefficients[1]) 

    return udf(calculate_slope_internal, DoubleType())


df = df.withColumn("slope", calculate_slope("array_1", "array_2")(df))

错误原因 #

1. Spark ML组件无法序列化：LinearRegression和VectorAssembler包含_thread.RLock等无法序列化的对象，UDF需要将所有依赖序列化后发送到Executor执行，因此触发序列化失败。
2. Executor端创建DataFrame是反模式：UDF内部调用spark.createDataFrame会尝试在Executor端使用Spark Session，而Session对象无法跨节点正常工作，必然导致分布式执行异常。

解决方案 #

放弃使用Spark ML组件，直接用单变量线性回归的数学公式计算斜率。公式如下：

slope = (n*Σ(xy) - Σx*Σy) / (n*Σ(x²) - (Σx)²)

其中：

• n为数组长度
• Σx是array_1所有元素的和
• Σy是array_2所有元素的和
• Σ(xy)是对应元素乘积的和
• Σ(x²)是array_1元素平方的和

这种方式不需要依赖Spark ML的分布式训练组件，完全适配UDF的序列化要求，同时性能更优。

实现代码 #

from pyspark.sql import functions as F
from pyspark.sql.types import DoubleType

# 构造测试数据
data = [(1, [1, 2, 3, 4], [8, 9, 7, 3]), (2, [5, 6, 7, 8], [1, 2, 3, 4]), (3, [9, 10, 11, 12], [1, 2, 3, 4])]
df = spark.createDataFrame(data, ["id", "array_1", "array_2"])

# 定义计算斜率的UDF，包含异常处理
def compute_slope(x_array, y_array):
    n = len(x_array)
    # 处理空数组或长度为0的情况
    if n == 0:
        return 0.0
    sum_x = sum(x_array)
    sum_y = sum(y_array)
    sum_xy = sum(x * y for x, y in zip(x_array, y_array))
    sum_x_sq = sum(x ** 2 for x in x_array)
    
    denominator = n * sum_x_sq - sum_x ** 2
    # 避免除以0（比如x数组所有元素相同的情况）
    if denominator == 0:
        return 0.0
    slope = (n * sum_xy - sum_x * sum_y) / denominator
    return slope

# 注册UDF
slope_udf = F.udf(compute_slope, DoubleType())

# 新增Slope列
result_df = df.withColumn("Slope", slope_udf(F.col("array_1"), F.col("array_2")))

# 查看结果
result_df.show()

运行结果 #

+---+------------+------------+-------------------+
| id|     array_1|     array_2|              Slope|
+---+------------+------------+-------------------+
|  1|[1, 2, 3, 4]|[8, 9, 7, 3]| -1.5999999999999996|
|  2|[5, 6, 7, 8]|[1, 2, 3, 4]|                1.0|
|  3|[9,10,11,12]|[1, 2, 3, 4]|                1.0|
+---+------------+------------+-------------------+

补充说明 #

• 加入了空数组、分母为0的异常处理，避免运行时抛出错误。
• 该方案无需为每行创建临时DataFrame和训练模型，性能远优于原方案，减少了大量不必要的分布式开销。

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