pyspark如何根据值添加选定的列

Posted 2023-04-15

【中文标题】pyspark如何根据值添加选定的列

【英文标题】：pyspark how to add selected columns based on value

【发布时间】：2021-04-07 09:44:44

【问题描述】：

对于下面的数据结构，我希望返回一个基于condition 列的新数据框。例如，如果"condition" =='A' 新数据帧应该在 group1 中有 cols 值，如果"condition" =='B' 新数据帧应该在 group2 中有 cols 值。问题是我不想硬编码列名，因为anothervalue 之后可能有很多列。我怎么能这样做？非常感谢您的帮助。例如对于这个输入数据框，

+---------+---------+---------+
|condition|   group1|   group2|
+---------+---------+---------+
|        A|SEA, WA|PDX, OR|
|        B| NY, NY| LA, CA|
+---------+---------+---------+

我想得到这个输出：

+---------+---------+
|condition|   group |  
+---------+---------+
|        A|SEA, WA|
|        B| LA, CA|
+---------+---------+

上面的输入数据框是使用这个 json 模式创建的：

jsonStrings = ['"condition":"A","group1":"city":"SEA","state":"WA","group2":"city":"PDX","state":"OR"','"condition":"B","group1":"city":"NY","state":"NY","group2":"city":"LA","state":"CA"']

【问题讨论】：

任何示例数据框？

我编辑了这个问题。谢谢

我还不清楚。能否以数据帧格式而不是 json 格式显示，能否也提供预期的输出数据帧？

我用示例数据框和 json 输入编辑了这个问题。

为什么不直接使用when？

【参考方案1】：

您可以简单地使用when 并构造如下的动态条件列表

from pyspark.sql.functions import *

conditions = when(col('condition') == 'A', col("group1"))\
    .when(col('condition') == 'B', col("group2")).otherwise(None)

df1.select(col('condition'), conditions.alias("group")).show(truncate=False)

输出：

+---------+---------+
|condition|group    |
+---------+---------+
|A        |SEA, WA|
|B        |LA, CA |
+---------+---------+