首页 > 解决方案 > Spark dataframe groupby 列中的唯一值

问题描述

我有以下数据框:

val df = Seq(
    ("A", 2.0),
    ("A", 1.5),
    ("B", 8.0),
    ("B", 9.0)
  ).toDF("id", "val")

我想做一个组,每个组基于唯一的ID 1.有一个运行计数,即第一个是0,第二个是1 2.组内的总计数。

应该看起来像

+---+---+---+---+
| id|val|order|count|
+---+---+---+---+
|  A|2.0|0 |2
|  A|1.5|1 |2
|  A|2.5|2 |2
|  B|8.0|0|2
|  B|9.0|1 |2
+---+---+

我看不到如何使用 spark sql 或函数执行此操作。

标签: apache-spark-sql

解决方案


这是一种方法。

输入数据:

+---+---+
|id |val|
+---+---+
|A  |2.0|
|A  |1.5|
|A  |4.5|
|A  |0.5|
|B  |8.0|
|B  |9.0|
+---+---+

使用 row_number 函数获取订单计数。

val w1 = Window.partitionBy("id").orderBy(lit(null))

df.withColumn("rank", row_number().over(w1))
  .withColumn("order", 'rank - 1)
  .withColumn("count",count('order).over(w1))
  .drop('rank)
  .orderBy('id)
  .show(false)

两者都给出相同的结果:

+---+---+-----+-----+
|id |val|order|count|
+---+---+-----+-----+
|A  |2.0|0    |4    |
|A  |1.5|1    |4    |
|A  |4.5|2    |4    |
|A  |0.5|3    |4    |
|B  |8.0|0    |2    |
|B  |9.0|1    |2    |
+---+---+-----+-----+

推荐阅读