Pyspark Dataframe将函数应用于两列

时间:2016-11-02 20:42:22

标签: pyspark spark-dataframe pyspark-sql

假设我有两个PySpark DataFrame df1df2

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

永远运行而不会给出错误/输出。

我的问题是:

  1. 由于我是Spark的新手,我构建输出DataFrame的方法是否有效? (我的方法是首先为所有ab值创建一个距离矩阵,然后找到min一个
  2. 我的代码的最后一行有什么问题以及如何修复它?

1 个答案:

答案 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|
+---+---+