Pyspark UDF pickle错误，无法pickle SwigPyObject对象

Question

我正在尝试将udf函数应用于由字符串组成的dataframe列。函数使用Tensorflow GUSE并将字符串转换为浮点数组。

import tensorflow as tf
import tensorflow_hub as hub
import numpy as np
import tf_sentencepiece
# Graph set up.
g = tf.Graph()
with g.as_default():
  text_input = tf.placeholder(dtype=tf.string, shape=[None])
  embed = hub.Module("https://tfhub.dev/google/universal-sentence-encoder-multilingual-large/1")
  embedded_text = embed(text_input)
  init_op = tf.group([tf.global_variables_initializer(), tf.tables_initializer()])
g.finalize()

# Initialize session.
session = tf.Session(graph=g)
session.run(init_op)

def embed_mail(x): 
    embedding = session.run(embedded_text, feed_dict={text_input:[x]})
    embedding = flatten(embedding)
    result = [np.float32(i).item() for i in embedding]
    return result

但每当我尝试使用以下命令运行此函数时：

embed_mail_udf = udf(embed_mail, ArrayType(FloatType()))
df = df.withColumn('embedding',embed_mail_udf(df.text))

我一直收到一个错误:无法序列化对象: TypeError: can't pickle SwigPyObject objects。我做错了什么？

Answer 1

要在集群上运行UDF的代码，Spark需要能够序列化“附加”到该函数的所有数据。您的tf.Session包含对TF embed_mail的引用，因此该函数是一个closure，而Spark首先需要序列化tf.Session的内容。我敢打赌，这是问题的一个原因。不幸的是，我没有使用TF的经验，但似乎你可以在运行Spark之前从TF获取邮件数据，广播它，然后在你的udf中使用？