Spark: как сделать значение нового столбца на основе разных столбцов

Question

Spark 2.2.1 Pyspark

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"])

А затем у меня есть список полей, которые мне интересны выше, из столбца "Pet"

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

Я хочу написать коддаст мне значение от dog_30, cat_30 или parrot_30, которое соответствует значению в "питомце".Например, в первой строке значение для столбца pet равно dog, и поэтому мы берем значение для dog_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   |
+--------+------+------+---------+------+

Answer 1

логика выключена;вам нужно .when(F.col("pet") == p, F.col(p + '_30')):

mycols = [F.when(F.col("pet") == p, F.col(p + '_30')) for p in dfvalues]
df = df.withColumn("newCol2",F.coalesce(F.coalesce(*mycols),F.lit(0)))
df.show()
+--------+------+------+---------+-------+
|     pet|dog_30|cat_30|parrot_30|newCol2|
+--------+------+------+---------+-------+
|     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|
+--------+------+------+---------+-------+