scala - Spark 数据帧 : Set column values if an conditional row is encountered

Question

我需要根据其他行的值条件来操作一行的值。

场景:如果任何行具有 (KEY=111 & IND=Yes)，则需要将 KEY 值设置为“999”，只有具有 KEY=202。如果这一行不存在，那么我们不应该设置任何东西。

示例输入:

val rowvalues = List(List("111", "Yes", "2016-11-15T00:00:00.000"), List("111", "No", "2016-12-01T00:00:00.000"),
  List("202", "NA", "2016-12-01T00:00:00.000"), List("202", "NA", "2016-11-15T00:00:00.000"),
  List("303", "Yes", "2016-12-01T00:00:00.000"), List("303", "NA", "2016-11-15T00:00:00.000"))
  .map(row => (row(0), row(1), row(2)))

val df = rowvalues.toDF("KEY", "IND", "DATE")

df.show(false)

+---+---+-----------------------+
|KEY|IND|DATE                   |
+---+---+-----------------------+
|111|Yes|2016-11-15T00:00:00.000|-----> Condition met as KEY=111 & IND=Yes
|111|No |2016-12-01T00:00:00.000|
|202|NA |2016-12-01T00:00:00.000|-----> Set KEY=999
|202|NA |2016-11-15T00:00:00.000|-----> Set KEY=999
|303|Yes|2016-12-01T00:00:00.000|
|303|NA |2016-11-15T00:00:00.000|
+---+---+-----------------------+

预期输出:

+---+---+-----------------------+
|KEY|IND|DATE                   |
+---+---+-----------------------+
|111|Yes|2016-11-15T00:00:00.000|
|111|No |2016-12-01T00:00:00.000|
|999|NA |2016-12-01T00:00:00.000|*
|999|NA |2016-11-15T00:00:00.000|*
|303|Yes|2016-12-01T00:00:00.000|
|303|NA |2016-11-15T00:00:00.000|
+---+---+-----------------------+

Answer 1

计算符合给定条件的行数并使用 if-else 语句。 when 和 otherwise 可以在将 202 替换为 999 时使用。

val cond = df.filter($"KEY" === 111  && $"IND" === "Yes").count() > 0

val df2 = if(cond) {
  df.withColumn("KEY", when($"KEY" === 202, 999).otherwise($"KEY"))
} else {
  df
}

结果:

+---+---+-----------------------+
|KEY|IND|DATE                   |
+---+---+-----------------------+
|111|Yes|2016-11-15T00:00:00.000|
|111|No |2016-12-01T00:00:00.000|
|999|NA |2016-12-01T00:00:00.000|
|999|NA |2016-11-15T00:00:00.000|
|303|Yes|2016-12-01T00:00:00.000|
|303|NA |2016-11-15T00:00:00.000|
+---+---+-----------------------+