scala - Spark Dataframe Group 通过新的指标列

Question

我需要按“KEY”列进行分组，并需要检查“TYPE_CODE”列是否同时具有“PL”和“JL”值，如果是这样，那么我需要添加一个指示符列作为“Y”否则“N”

示例:

    //Input Values
    val values = List(List("66","PL") ,
    List("67","JL") , List("67","PL"),List("67","PO"),
    List("68","JL"),List("68","PO")).map(x =>(x(0), x(1)))

    import spark.implicits._
    //created a dataframe
    val cmc = values.toDF("KEY","TYPE_CODE")

    cmc.show(false)
    ------------------------
    KEY |TYPE_CODE  |
    ------------------------
    66  |PL |
    67  |JL |
    67  |PL |
    67  |PO |
    68  |JL |
    68  |PO |
    -------------------------

预期输出:

对于每个“KEY”，如果它具有“TYPE_CODE”同时具有 PL 和 JL 则 Y否则N

    -----------------------------------------------------
    KEY |TYPE_CODE  | Indicator
    -----------------------------------------------------
    66  |PL         | N
    67  |JL         | Y
    67  |PL         | Y
    67  |PO         | Y
    68  |JL         | N
    68  |PO         | N
    ---------------------------------------------------

例如，67 同时具有 PL 和 JL - 所以“Y”66 只有 PL - 所以“N”68 只有 JL - 所以“N”

Answer 1

一个选项:

1) 收集TYPE_CODE作为列表；

2) 检查是否包含特定字符串；

3) 然后使用explode压平列表:

(cmc.groupBy("KEY")
    .agg(collect_list("TYPE_CODE").as("TYPE_CODE"))
    .withColumn("Indicator", 
        when(array_contains($"TYPE_CODE", "PL") && array_contains($"TYPE_CODE", "JL"), "Y").otherwise("N"))
    .withColumn("TYPE_CODE", explode($"TYPE_CODE"))).show
+---+---------+---------+
|KEY|TYPE_CODE|Indicator|
+---+---------+---------+
| 68|       JL|        N|
| 68|       PO|        N|    
| 67|       JL|        Y|
| 67|       PL|        Y|
| 67|       PO|        Y|
| 66|       PL|        N|
+---+---------+---------+