apache-spark - Spark : Set a column value based on multiple row conditions

Question

我有以下格式的数据框:

+----+---+-----+------+-----+------+
|AGEF|SEX|F0_34|F35_44|M0_34|M35_44|
+----+---+-----+------+-----+------+
|  30|  0|    0|     0|    0|     0|
|  94|  1|    0|     0|    0|     0|
|  94|  0|    0|     0|    0|     0|
|  94|  0|    0|     0|    0|     0|
|  94|  1|    0|     0|    0|     0|
|  44|  0|    0|     0|    0|     0|
|  66|  0|    0|     0|    0|     0|
|  66|  0|    0|     0|    0|     0|
|  74|  0|    0|     0|    0|     0|
|  74|  0|    0|     0|    0|     0|
|  29|  0|    0|     0|    0|     0|

现在根据列 AGEF 和 SEX 的值，我需要将 1 分配给相应的列名。每个列名称都是不言自明的，例如 F0_34 是 0 到 34 岁之间的女性，其他情况类似。

预期输出是

+----+---+-----+------+-----+------+
|AGEF|SEX|F0_34|F35_44|M0_34|M35_44|
+----+---+-----+------+-----+------+
|  30|  0|    1|     0|    0|     0|
|  94|  1|    0|     0|    0|     0|
|  94|  0|    0|     0|    0|     0|
|  94|  0|    0|     0|    0|     0|
|  94|  1|    0|     0|    0|     0|
|  44|  0|    0|     1|    0|     0|
|  66|  0|    0|     0|    0|     0|
|  66|  0|    0|     0|    0|     0|
|  74|  0|    0|     0|    0|     0|
|  74|  0|    0|     0|    0|     0|
|  29|  0|    1|     0|    0|     0|

提前致谢!!!

Answer 1

通常最有效的方法是直接对 SQL 表达式进行操作。例如:

def categorize(ageRanges: Seq[(Int, Int)], sexValues: Seq[(Int, String)]) = for {
  (ageL, ageH) <- ageRanges
  (sexV, sexL) <- sexValues
} yield ($"SEX" === sexL && $"AGEF".between(ageL, ageH)).alias(
  s"$sexL-$ageL-$ageH"
)

df.select(
  $"*" +: categorize(Seq((0, 34), (35, 44)), Seq((0, "F"), (1, "M"))): _*
)