根据 Spark Scala 中的特定值修改 DataFrame 值

Question

查看我的代码：

val spark: SparkSession = SparkSession.builder().appName("ReadFiles").master("local[*]").getOrCreate()

import spark.implicits._

val 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, lit(nums(0))).otherwise(col("newCol1")))
  .withColumn("newCol2", when(col("id") === 4, lit(nums(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|
//+---+----+----+-------+-------+