将不带返回值的Python Lambda函数转换为Pyspark

Question

我在Python语言中有一个有效的lambda函数，它可以计算dataset1中的每个字符串与dataset2中的字符串之间的最高相似度。在迭代期间，它将字符串、最佳匹配和相似度以及其他一些信息写入bigquery。没有返回值，因为该函数的目的是向bigquery数据集中插入一行。这个过程需要相当长的时间，这就是为什么我想使用Pyspark和Dataproc来加速这个过程。

将熊猫数据帧转换为spark很容易。我在注册udf时遇到了问题，因为它没有返回值，而pyspark需要一个返回值。此外，我不知道如何将python中的'apply‘函数映射到pyspark变体。因此，基本上我的问题是如何转换下面的python代码来处理spark dataframe。

以下代码可在常规Python环境中运行：

def embargomatch(name, code, embargo_names):
     find best match 
     insert best match and additional information to bigquery

customer_names.apply(lambda x: embargoMatch(x['name'], x['customer_code'],embargo_names),axis=1)

因为pyspark需要返回类型，所以我在udf中添加了'return 1‘，并尝试了以下操作：

customer_names = spark.createDataFrame(customer_names)

from pyspark.sql.types import IntegerType
embargo_match_udf = udf(lambda x: embargoMatch(x['name'], x['customer_code'],embargo_names), IntegerType())

现在，我在尝试应用select函数时遇到了困难，因为我不知道要给出什么参数。

Answer 1

我怀疑您遇到了如何将多列传递给udf的问题--对于这个问题，有一个很好的答案：Pyspark: Pass multiple columns in UDF。

与其基于封装您的函数的lambda创建udf，不如考虑通过直接创建基于embargomatch的udf来简化。

embargo_names = ...

# The parameters here are the columns passed into the udf
def embargomatch(name, customer_code):
    pass
embargo_match_udf = udf(embargomatch, IntegerType())
customer_names.select(embargo_match_udf(array('name', 'customer_code')).alias('column_name'))

也就是说，您的udf不会返回任何内容--我通常认为udf是一种向数据帧中添加列的方法，但没有副作用。如果你想在bigquery中插入记录，可以考虑这样做：

customer_names.select('column_name').write.parquet('gs://some/path')
os.system("bq load --source_format=PARQUET [DATASET].[TABLE] gs://some/path")