Spark:如何根据不同的列生成新列的值

火花2.2.1

df = sqlContext.createDataFrame([
    ("dog", "1", "2", "3"),
    ("cat", "4", "5", "6"),
    ("dog", "7", "8", "9"),
    ("cat", "10", "11", "12"),
    ("dog", "13", "14", "15"),
    ("parrot", "16", "17", "18"),
    ("goldfish", "19", "20", "21"),
], ["pet", "dog_30", "cat_30", "parrot_30"])

dfvalues = ["dog", "cat", "parrot"]

我想写的代码塔特会给我的价值从 dog_30 , cat_30 或 parrot_30 对应于“pet”中的值。例如,在第一行中 pet 列为 dog 狗\u 30 这是1。

我试着用它来获取代码,但它只给了我空的列 stats . 我也不知道怎么处理这个问题 goldfish 案例。我想把它设为0。

mycols = [F.when(F.col("pet") == p + "_30", p) for p in dfvalues]
df = df.withColumn("newCol2",F.coalesce(*stats) )
df.show()

期望输出:

+--------+------+------+---------+------+
|     pet|dog_30|cat_30|parrot_30|stats |
+--------+------+------+---------+------+
|     dog|     1|     2|        3|  1   |
|     cat|     4|     5|        6|  5   |
|     dog|     7|     8|        9|  7   |
|     cat|    10|    11|       12|  11  |
|     dog|    13|    14|       15|  13  |
|  parrot|    16|    17|       18|  18  |
|goldfish|    19|    20|       21|  0   |
+--------+------+------+---------+------+