Добавление нового столбца на основе старого столбца в Spark DataFrame

Question

У меня есть следующий фрейм данных.

key     | value
inv_1_c | 5
inv_1_v | 8
inv_2_c | 9

Я хотел бы добавить два столбца к кадру данных Напряжение и ток.

Voltage будет значением, если ключ заканчивается на "_v" или 0 в противном случае.

Current будет значением, если ключ заканчивается на "_c" или 0 в противном случае.

Каким будет для этого искровой код scala?

Answer 1

Вы можете использовать функцию subString, чтобы получить последние два символа и проверить, содержит ли она _v или _c, и добавить два новых столбца с withColumn

import org.apache.spark.sql.functions._

val data = Seq(
  ("inv_1_c", "5"),
  ("inv_1_v", "8"),
  ("inv_2_c", "9")
).toDF("key", "value")

data.withColumn("temp", substring($"key", -2, 2))
    .withColumn("voltage", when($"temp" === "_v", $"value").otherwise(0))
    .withColumn("current", when($"temp" === "_c", $"value").otherwise(0))
    .drop("temp")

Выход:

+-------+-----+-------+-------+
|key    |value|voltage|current|
+-------+-----+-------+-------+
|inv_1_c|5    |0      |5      |
|inv_1_v|8    |8      |0      |
|inv_2_c|9    |0      |9      |
+-------+-----+-------+-------+

Надеюсь, это поможет!