Изменить значения DataFrame для определенного значения в Spark Scala

Question

См. Мой код:

  val spark: SparkSession = SparkSession.builder().appName("ReadFiles").master("local[*]").getOrCreate()
  val sc: SparkContext = spark.sparkContext
  import spark.implicits._

  var data: DataFrame = spark.read.option("header", "true")
    .option("inferschema", "true")
    .csv("Resources/atom.csv")
     data.show()

Данные выглядят следующим образом:

ID  Name  City  newCol1  newCol2
1   Ali   lhr   null     null
2   Ahad  swl    1        10
3   Sana  khi   null     null
4   ABC   xyz   null     null

Новый список значений:

val nums: List[Int] = List(10,20)

Я хочу добавить эти значения где ID=4. Так что DataFrame может выглядеть так:

ID  Name  City  newCol1  newCol2
1   Ali   lhr   null     null
2   Ahad  swl    1        10
3   Sana  khi   null     null
4   ABC   xyz    10       20

Интересно, возможно ли это? Любая помощь будет оценена. Спасибо

Answer 1

Возможно, используйте when otherwise заявления для этого случая.

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

df.withColumn("newCol1",when(col("id") === 4,split(lit(nums.mkString(",")),",")(0)).otherwise(col("newCol1"))).
withColumn("newCol2",when(col("id") === 4,split(lit(nums.mkString(",")),",")(1)).otherwise(col("newCol2"))).
show()
//+---+----+----+-------+-------+
//| ID|Name|City|newCol1|newCol2|
//+---+----+----+-------+-------+
//|  1| Ali| lhr|   null|   null|
//|  2|Ahad| swl|      1|     10|
//|  3|Sana| khi|   null|   null|
//|  4| ABC| xyz|     10|     20|
//+---+----+----+-------+-------+