根据其他两列的最小值添加新列

5fjcxozz  于 2021-05-29  发布在  Spark
关注(0)|答案(1)|浏览(464)

我有一个如下所示的Dataframe

id  id1 seq_nbr     id2 orig_id1  orig_id2  
1   740     2       217     740     217         
1   740     3       216     740     216
1   740     4       216     740     216
1   740     5       217     740     217
1   367     1       217     740     217

我需要在orig\u id1,orig\u id2的基础上添加名为orig\u seq\u nbr的新列。预期输出为

id  id1 seq_nbr     id2 orig_id1  orig_id2 orig_seq_nbr
1   740     3       216     740     216         3
1   740     4       216     740     216         3
1   740     5       217     740     217         2
1   367     1       217     740     217         2
1   740     2       217     740     217         2

前两行有3,因为orig_id1,orig_id2有740216,seq_nbr的min基于id1,id2列(740216)是3。最后三行有2,因为orig_id1,orig_id2有740217,seq_nbr基于id1,id2列(740217)的最小值是2。
我试图根据orig\u id1,orig\u id2找到min,但得到了下面的结果

val win = Window.partitionBy("orig_id1", "orig_id2")
df1.withColumn("orig_seq_nbr",min(col("seq_nbr")).over(win)).show()
id  id1 seq_nbr     id2 orig_id1  orig_id2 orig_seq_nbr
1   740     3       216     740     216         3
1   740     4       216     740     216         3
1   740     5       217     740     217         1
1   367     1       217     740     217         1
1   740     2       217     740     217         1

我应该做什么更改才能将最后三行的orig\ u seq\ u nbr设置为2?

yquaqz18

yquaqz181#

试试这个方法-

1. 加载测试数据

val data =
      """
        |id | id1 | seq_nbr   |  id2 |orig_id1 | orig_id2
        |1  | 740 |     2     |  217 |    740  |   217
        |1  | 740 |     3     |  216 |    740  |   216
        |1  | 740 |     4     |  216 |    740  |   216
        |1  | 740 |     5     |  217 |    740  |   217
        |1  | 367 |     1     |  217 |    740  |   217
      """.stripMargin

    val stringDS1 = data.split(System.lineSeparator())
      .map(_.split("\\|").map(_.replaceAll("""^[ \t]+|[ \t]+$""", "")).mkString(","))
      .toSeq.toDS()
    val df1 = spark.read
      .option("sep", ",")
      .option("inferSchema", "true")
      .option("header", "true")
      .option("nullValue", "null")
      .csv(stringDS1)
    df1.show(false)
    df1.printSchema()

    /**
      * +---+---+-------+---+--------+--------+
      * |id |id1|seq_nbr|id2|orig_id1|orig_id2|
      * +---+---+-------+---+--------+--------+
      * |1  |740|2      |217|740     |217     |
      * |1  |740|3      |216|740     |216     |
      * |1  |740|4      |216|740     |216     |
      * |1  |740|5      |217|740     |217     |
      * |1  |367|1      |217|740     |217     |
      * +---+---+-------+---+--------+--------+
      *
      * root
      * |-- id: integer (nullable = true)
      * |-- id1: integer (nullable = true)
      * |-- seq_nbr: integer (nullable = true)
      * |-- id2: integer (nullable = true)
      * |-- orig_id1: integer (nullable = true)
      * |-- orig_id2: integer (nullable = true)
      */

2. 使用dsl

val win = Window.partitionBy("orig_id1", "orig_id2")
    df1.withColumn("orig_seq_nbr",
      min(when($"orig_id1" === $"id1" && $"orig_id2" === $"id2", $"seq_nbr").otherwise(null))
      .over(win)
    ).show()

    /**
      * +---+---+-------+---+--------+--------+------------+
      * | id|id1|seq_nbr|id2|orig_id1|orig_id2|orig_seq_nbr|
      * +---+---+-------+---+--------+--------+------------+
      * |  1|740|      3|216|     740|     216|           3|
      * |  1|740|      4|216|     740|     216|           3|
      * |  1|740|      2|217|     740|     217|           2|
      * |  1|740|      5|217|     740|     217|           2|
      * |  1|367|      1|217|     740|     217|           2|
      * +---+---+-------+---+--------+--------+------------+
      */

3. 使用sql表达式

df1.withColumn("orig_seq_nbr",
      expr("min(case when orig_id1=id1 and orig_id2=id2 then seq_nbr else NULL end) " +
        "OVER (PARTITION BY orig_id1, orig_id2) ")
    ).show()

    /**
      * +---+---+-------+---+--------+--------+------------+
      * | id|id1|seq_nbr|id2|orig_id1|orig_id2|orig_seq_nbr|
      * +---+---+-------+---+--------+--------+------------+
      * |  1|740|      3|216|     740|     216|           3|
      * |  1|740|      4|216|     740|     216|           3|
      * |  1|740|      2|217|     740|     217|           2|
      * |  1|740|      5|217|     740|     217|           2|
      * |  1|367|      1|217|     740|     217|           2|
      * +---+---+-------+---+--------+--------+------------+
      */

相关问题