如果两个列值在另一个数据帧中,如何在pyspark中添加列?

时间:2017-04-17 01:54:00

标签: python python-2.7 apache-spark pyspark spark-dataframe

我对pyspark很新。我有两个这样的数据帧:

DF1: enter image description here

DF2: enter image description here

df1中的

label列最初不存在。我后来加了。如果[user_id, sku_id]对df1在df2中,那么我想在df1中添加一个列并将其设置为1,否则为0,就像df1所示。我怎么能在pyspark做到这一点?我正在使用py2.7。

1 个答案:

答案 0 :(得分:0)

首先在两个数据帧上执行左外连接,然后在右数据帧的一列上使用whenotherwise函数。这是我试过的完整解决方案 -

from pyspark.sql import functions as F
from pyspark.sql.functions import col

# this is just data input
data1 = [[4,3,3],[2,4,3],[4,2,4],[4,3,3]]
data2 = [[4,3,3],[2,3,3],[4,1,4]]

# create dataframes
df1 = spark.createDataFrame(data1,schema=['userId','sku_id','type'])
df2 = spark.createDataFrame(data2,schema=['userId','sku_id','type'])

# condition for join
cond=[df1.userId==df2.userId,df1.sku_id==df2.sku_id,df1.type==df2.type]

# magic
df1.join(df2,cond,how='left_outer')\
  .select(df1.userId,df1.sku_id,df1.type,df2.userId.alias('uid'))\
  .withColumn('label',F.when(col('uid')>0  ,1).otherwise(0))\
  .drop(col('uid'))\
  .show()

输出:

+------+------+----+-----+
|userId|sku_id|type|label|
+------+------+----+-----+
|     2|     4|   3|    0|
|     4|     3|   3|    1|
|     4|     3|   3|    1|
|     4|     2|   4|    0|
+------+------+----+-----+