假设我有两个PySpark DataFrame df1
和df2
。
df1= 'a'
1
2
5
df2= 'b'
3
6
我希望为每个df2['b']
找到最接近的df1['a']
值,并将最接近的值添加为df1
中的新列。
换句话说,对于x
中的每个值df1['a']
,我想找到一个y
,为min(abx(x-y))
找到y in df2['b']
(注意:可以假设只有一个y
可以达到最小距离),结果将是
'a' 'b'
1 3
2 3
5 6
我尝试使用以下代码首先创建距离矩阵(在找到达到最小距离的值之前):
from pyspark.sql.types import IntegerType
from pyspark.sql.functions import udf
def dict(x,y):
return abs(x-y)
udf_dict = udf(dict, IntegerType())
sql_sc = SQLContext(sc)
udf_dict(df1.a, df2.b)
给出了
Column<PythonUDF#dist(a,b)>
然后我试了
sql_sc.CreateDataFrame(udf_dict(df1.a, df2.b))
永远运行而不会给出错误/输出。
我的问题是:
a
和b
值创建一个距离矩阵,然后找到min
一个答案 0 :(得分:7)
从第二个问题开始 - 您只能将udf应用于现有数据框,我认为您正在考虑这样的事情:
>>> df1.join(df2).withColumn('distance', udf_dict(df1.a, df2.b)).show()
+---+---+--------+
| a| b|distance|
+---+---+--------+
| 1| 3| 2|
| 1| 6| 5|
| 2| 3| 1|
| 2| 6| 4|
| 5| 3| 2|
| 5| 6| 1|
+---+---+--------+
但是使用内部abs
:
>>> from pyspark.sql.functions import abs
>>> df1.join(df2).withColumn('distance', abs(df1.a -df2.b))
然后你可以通过计算找到匹配的数字:
>>> distances = df1.join(df2).withColumn('distance', abs(df1.a -df2.b))
>>> min_distances = distances.groupBy('a').agg(min('distance').alias('distance'))
>>> distances.join(min_distances, ['a', 'distance']).select('a', 'b').show()
+---+---+
| a| b|
+---+---+
| 5| 6|
| 1| 3|
| 2| 3|
+---+---+