在UDF PySpark中传递多个列

Question

我想计算PySpark数据帧的两列之间的Jaro Winkler距离。Jaro-Winkler距离可通过所有节点上的pyjarowinkler包获得。

pyjarowinkler的工作原理如下：

from pyjarowinkler import distance
distance.get_jaro_distance("A", "A", winkler=True, scaling=0.1)

输出：

1.0

我试图编写一个UDF，将两列作为序列传递，并使用lambda函数计算距离。我是这样做的：

@pandas_udf("float", PandasUDFType.SCALAR)
def get_distance(col1, col2):
    import pandas as pd
    distance_df  = pd.DataFrame({'column_A': col1, 'column_B': col2})
    distance_df['distance'] = distance_df.apply(lambda x: distance.get_jaro_distance(str(distance_df['column_A']), str(distance_df['column_B']), winkler = True, scaling = 0.1))
    return distance_df['distance']

temp = temp.withColumn('jaro_distance', get_distance(temp.x, temp.x))

我应该能够在上述函数中传递任意两个字符串列。我得到以下输出：

+---+---+---+-------------+
|  x|  y|  z|jaro_distance|
+---+---+---+-------------+
|  A|  1|  2|         null|
|  B|  3|  4|         null|
|  C|  5|  6|         null|
|  D|  7|  8|         null|
+---+---+---+-------------+

预期产出：

+---+---+---+-------------+
|  x|  y|  z|jaro_distance|
+---+---+---+-------------+
|  A|  1|  2|          1.0|
|  B|  3|  4|          1.0|
|  C|  5|  6|          1.0|
|  D|  7|  8|          1.0|
+---+---+---+-------------+

我怀疑这可能是因为str（distance_df['column_A']）不正确。它包含所有行值的连接字符串。

虽然这个代码为我工作：

@pandas_udf("float", PandasUDFType.SCALAR)
def get_distance(col):
    return col.apply(lambda x: distance.get_jaro_distance(x, "A", winkler = True, scaling = 0.1))

temp = temp.withColumn('jaro_distance', get_distance(temp.x))

输出：

+---+---+---+-------------+
|  x|  y|  z|jaro_distance|
+---+---+---+-------------+
|  A|  1|  2|          1.0|
|  B|  3|  4|          0.0|
|  C|  5|  6|          0.0|
|  D|  7|  8|          0.0|
+---+---+---+-------------+

熊猫UDF有办法做到这一点吗？我正在处理数百万张唱片，所以UDF会很贵，但如果有效，仍然可以接受。谢谢。

Answer 1

您可以首先合并所有数据帧，在分区被洗牌并分发到工作节点之后，使用相同的分区键进行分区，并在计算之前恢复它们。请查看我为这个场景编写了一个小工具包的示例：SparkyPandas

Answer 2

错误来自df中的函数。应用方法，将其调整到以下位置，以便修复：

@pandas_udf("float", PandasUDFType.SCALAR)
def get_distance(col1, col2):
    import pandas as pd
    distance_df  = pd.DataFrame({'column_A': col1, 'column_B': col2})
    distance_df['distance'] = distance_df.apply(lambda x: distance.get_jaro_distance(x['column_A'], x['column_B'], winkler = True, scaling = 0.1), axis=1)
    return distance_df['distance']

然而，熊猫却不例外。apply方法不是矢量化的，这不符合我们在PySpark中需要pandas_udf而不是udf的目的。一个更快、开销更小的解决方案是使用列表理解来创建返回的pd。系列（有关Pandas df.apply及其备选方案的更多讨论，请查看此链接）：

from pandas import Series

@pandas_udf("float", PandasUDFType.SCALAR)
def get_distance(col1, col2):
   return Series([ distance.get_jaro_distance(c1, c2, winkler=True, scaling=0.1) for c1,c2 in zip(col1, col2) ])

df.withColumn('jaro_distance', get_distance('x', 'y')).show()
+---+---+---+-------------+
|  x|  y|  z|jaro_distance|
+---+---+---+-------------+
| AB| 1B|  2|         0.67|
| BB| BB|  4|          1.0|
| CB| 5D|  6|          0.0|
| DB|B7F|  8|         0.61|
+---+---+---+-------------+