在 lambda 中使用 pyspark sql 函数导致 pickle 错误

Question

我有一个如下图所示的pyspark代码：

if has_column(df_s_a, "x"):
    df6 = df_s_a \
        .withColumn("x", explode_outer(f.col("x"))) \
        .withColumn("y", 
                    f.udf(
                    lambda x: element_at("z.r.id", -1) if x=='C' else lit(None).cast("string"),
                        StringType()
                    )(col('x.type'))
                   ) \
            .select(
                from_unixtime(col("x.p")).alias("xp"),
                f.col("y"),
        )

我尝试在 lambda 内部对数据帧 d_s_a 进行任何访问，都会导致以下错误：

TypeError                                 Traceback (most recent call last)
/usr/local/lib/python3.9/site-packages/pyspark/serializers.py in dumps(self, obj)
    436         try:
--> 437             return cloudpickle.dumps(obj, pickle_protocol)
    438         except pickle.PickleError:
.
.
.
PicklingError: Could not serialize object: TypeError: cannot pickle '_thread.RLock' object

有没有一种方法可以在 lambda 中调用 SQL 函数？我也使用了 when() 子句，但它会评估 when & else 不管条件：Case when statement to conditionally populate columns in pyspark

Answer 1

您不能在 udf 中使用 Column 函数。 udf 从单行而不是整列中获取传递的元素。 因此，不要使用 lit(None)，而是使用 None。 而不是 element_at("z.r.id", -1) 使用 id[-1] 将 z.r.id 传递给 udf。