Using a pyspark sql function inside lambda causing pickle error

user_11697394 0
Viewed 262

I have a pyspark code as shown below:

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"),
        )

Any access i try to make inside the lambda to the dataframe d_s_a, it causes the below error:

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

Is there a way I can make that function call to SQL function within the lambda? I've this using when() clause as well, but it evaluates the when & otherwise regardless of the condition: Case when statement to conditionally populate columns in pyspark

pyspark apache-spark-sql apache-spark
1 Answers

You cannot use Column functions inside a udf. The udf gets passed elements from a single row, not an entire column. So instead of lit(None) just use None. And instead of element_at("z.r.id", -1) use id[-1] where you pass z.r.id to the udf.

user_165984700
Related