根据其他列上的where条件在Pyspark数据帧中添加新列

Question

我有一个Pyspark数据帧，如下所示：

+------------+-------------+--------------------+
|package_id  | location    | package_scan_code  | 
+------------+-------------+--------------------+
|123         | Denver      |05                  |  
|123         | LosAngeles  |03                  |  
|123         | Dallas      |09                  |  
|123         | Vail        |02                  | 
|456         | Jacksonville|05                  |  
|456         | Nashville   |09                  |
|456         | Memphis     |03                  |

“包扫描码”03代表包裹的来源。在

我想在这个数据帧中添加一列“origin”，这样对于每个包（由“package_id”标识），新添加的origin列中的值将与“package_scan_code”03对应。在

在上述情况下，有两个唯一的包123和456，它们分别来自洛杉矶和孟菲斯（对应包的扫描代码03）。在

所以我希望我的输出如下：

^{pr2}$

如何在Pyspark中实现这一点？我尝试了.withColumn方法，但是我不能正确地处理条件。在

Answer 1

不管数据帧中的每个package_id发生package_scan_code=03多少次，此代码都应该可以工作。我又添加了一个(123,'LosAngeles','03')来演示-

步骤1:创建数据帧

values = [(123,'Denver','05'),(123,'LosAngeles','03'),(123,'Dallas','09'),(123,'Vail','02'),(123,'LosAngeles','03'),
          (456,'Jacksonville','05'),(456,'Nashville','09'),(456,'Memphis','03')]
df = sqlContext.createDataFrame(values,['package_id','location','package_scan_code'])

第2步：创建package_id和location的字典。在

^{pr2}$

第3步：创建列，映射字典。在

from pyspark.sql.functions import col, create_map, lit
from itertools import chain
mapping_expr = create_map([lit(x) for x in chain(*dict_location_scan_code.items())])
df = df.withColumn('origin', mapping_expr.getItem(col('package_id')))
df.show()
+     +      +        -+     +
|package_id|    location|package_scan_code|    origin|
+     +      +        -+     +
|       123|      Denver|               05|LosAngeles|
|       123|  LosAngeles|               03|LosAngeles|
|       123|      Dallas|               09|LosAngeles|
|       123|        Vail|               02|LosAngeles|
|       123|  LosAngeles|               03|LosAngeles|
|       456|Jacksonville|               05|   Memphis|
|       456|   Nashville|               09|   Memphis|
|       456|     Memphis|               03|   Memphis|
+     +      +        -+     +

Answer 2

按package_scan_code == '03'筛选数据帧，然后与原始数据帧连接：

(df.filter(df.package_scan_code == '03')
   .selectExpr('package_id', 'location as origin')
   .join(df, ['package_id'], how='right')
   .show())
+     +     +      +        -+
|package_id|    origin|    location|package_scan_code|
+     +     +      +        -+
|       123|LosAngeles|      Denver|               05|
|       123|LosAngeles|  LosAngeles|               03|
|       123|LosAngeles|      Dallas|               09|
|       123|LosAngeles|        Vail|               02|
|       456|   Memphis|Jacksonville|               05|
|       456|   Memphis|   Nashville|               09|
|       456|   Memphis|     Memphis|               03|
+     +     +      +        -+

注意：这假设您最多有一个package_scan_code等于03，否则逻辑将不正确，您需要重新考虑如何定义origin。在