pyspark 如何根據值添加選定的列

Question

對於下面的數據結構，我希望根據condition列返回一個新的 dataframe。 例如，如果"condition" =='A' ，則新的 dataframe 應該在 group1 中具有 cols 值，如果"condition" =='B' ，則新的 dataframe 應該在 group2 中具有 cols 值。 問題是我不想對列名進行硬編碼，因為在anothervalue之后可能有很多列。 我怎么能這樣做？ 非常感謝您的幫助。 例如對於這個輸入 dataframe，

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

我想要這個 output：

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

上述輸入 dataframe 是使用此 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"}}']

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

Output：

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