比较DataFrames以获取布尔列

时间:2016-12-20 18:15:21

标签: apache-spark pyspark

我有和n x m DataFrame和1 x m DataFrame

df1=sc.parallelize([('a1',.5,.27),('a2',.15,.40),('a3',.7,.05)]).toDF(['id1','x1', 'x2'])

+---+----+----+
|id1|  x1|  x2|
+---+----+----+
| a1| 0.5|0.27|
| a2|0.15| 0.4|
| a3| 0.7|0.05|
+---+----+----+

df2=sc.parallelize([(.4,.3)]).toDF(['w1','w2'])


+---+---+
| w1| w2|
+---+---+
|0.4|0.3|
+---+---+

我想执行一个布尔操作,将df1中的列x1与df2中的列w1和df2中的列x2与df2中的列w2进行比较。我希望结果返回一个DataFrame,其中第一列是' id1'在df1。

我希望我的结果看起来像这样

+---+---+---+
|id1| x1| x2|
+---+---+---+
| a1|  1|  0|
| a2|  0|  1|
| a3|  1|  0|
+---+---+---+

我现在所拥有的只是

rd=df1.rdd
rd_list=df2.rdd.collect()


def function_1(x):
    bool_1=int(x[1]>rd_list[0][0])
    bool_2=int(x[2]>rd_list[0][1])
    return (x[0],bool_1,bool_2)
rd.map(function_1).toDF().show()

+---+---+---+
|id1| x1| x2|
+---+---+---+
| a1|  1|  0|
| a2|  0|  1|
| a3|  1|  0|
+---+---+---+

这让我得到了我的结果,但必须有更好的方法。

1 个答案:

答案 0 :(得分:2)

您可以与文字进行比较(使用单行数据框没有多大意义):

from pyspark.sql.functions import col

w1, w2 = df2.first()

df1.select(
    "id1", 
    (col("x1") > w1).cast("integer"),
    (col("x2") > w2).cast("integer")
).toDF("id1", "w1", "w2")

或应用交叉产品和select

from pyspark.sql.functions import broadcast

df1.crossJoin(broadcast(df2)).select(
    "id1",
    (col("x1") > col("w1")).cast("integer"), 
    (col("x2") > col("w2")).cast("integer")
).toDF("id1", "x1", "x2")

如果您使用Spark 2.0或更早版本,则join而不是crossJoinbroadcast结合后应具有相同的效果。