Spark Dataframe Group具有新指标列

时间:2017-09-12 14:08:39

标签: scala apache-spark dataframe apache-spark-sql

我需要按“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”

2 个答案:

答案 0 :(得分:4)

另一种选择:

  1. KEY分组并使用agg创建两个单独的指标列(一个用于JL,另一个用于PL),然后计算组合指标< / p>

  2. join使用原始DataFrame

  3. 共:

    val indicators = cmc.groupBy("KEY").agg(
      sum(when($"TYPE_CODE" === "PL", 1).otherwise(0)) as "pls",
      sum(when($"TYPE_CODE" === "JL", 1).otherwise(0)) as "jls"
    ).withColumn("Indicator", when($"pls" > 0 && $"jls" > 0, "Y").otherwise("N"))
    
    val result = cmc.join(indicators, "KEY")
      .select("KEY", "TYPE_CODE", "Indicator")
    

    这可能比@ Psidom的答案慢,但可能更安全 - collect_list如果你有一个特定键的大量匹配可能会有问题(该列表必须存储在单个工作人员的记忆中)。

    修改

    如果已知输入唯一(即JL / PL最多只出现一次,最多),可以使用简单的indicators创建count聚合,(可以说)更容易阅读:

    val indicators = cmc
      .where($"TYPE_CODE".isin("PL", "JL"))
      .groupBy("KEY").count()
      .withColumn("Indicator", when($"count" === 2, "Y").otherwise("N"))
    

答案 1 :(得分:2)

一个选项:

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|
+---+---------+---------+