Фреймы данных Pyspark: извлечение столбца на основе значения другого столбца

Question

У меня есть фрейм данных со следующими столбцами и соответствующими значениями (простите мое форматирование, но я не знаю, как поместить его в формат таблицы):

Src_ip     dst_ip     V1     V2     V3     top
"A"         "B"       xx     yy     zz     "V1"

Теперь я хочу добавить столбец, скажем top_value, который принимает значение столбца, соответствующего строке в V1.

Src_ip     dst_ip     V1     V2     V3     top   top_value
"A"         "B"       xx     yy     zz     "V1"     xx

Таким образом, в принципе, получите значение, соответствующее значению в столбце "top", и создайте новый столбец с именем "top_value"

Я пытался создать UDF, а также использоватьстрока как псевдоним, но не может этого сделать.Может кто-нибудь, пожалуйста, помогите.

Answer 1

Вы можете собрать столбцы V1, V2 и V3 как struct и передать функцию udf с помощью столбца top и извлечь значение как

scala

import org.apache.spark.sql.functions._
def findValueUdf = udf((strct: Row, top: String) => strct.getAs[String](top))

df.withColumn("top_value", findValueUdf(struct("V1", "V2", "V3"), col("top")))

, который должен дать вам

+------+------+---+---+---+---+---------+
|Src_ip|dst_ip|V1 |V2 |V3 |top|top_value|
+------+------+---+---+---+---+---------+
|A     |B     |xx |yy |zz |V1 |xx       |
+------+------+---+---+---+---+---------+

pyspark

эквивалентный код в pyspark будет

from pyspark.sql import functions as f
from pyspark.sql import types as t
def findValueUdf(strct, top):
    return strct[top]

FVUdf = f.udf(findValueUdf, t.StringType())

df.withColumn("top_value", FVUdf(f.struct("V1", "V2", "V3"), f.col("top")))

Более того, вы можете определить имена столбцовв списке, который будет использоваться в функции struct, чтобы вам не приходилось жестко их кодировать.

Надеюсь, ответ будет полезным