在 Spark Dataframe 的 window 上创建组 ID

Question

我有一个数据框，我想在其中为每个 Window 分区提供 ID。例如我有

id | col |
1  |  a  |
2  |  a  |
3  |  b  |
4  |  c  |
5  |  c  |

所以我想要（基于列 col 的分组）

id | group |
1  |  1    |
2  |  1    |
3  |  2    |
4  |  3    |
5  |  3    |

我想使用 window 函数，但无论如何我找不到为每个 window 分配一个 Id。我需要这样的东西：

w = Window().partitionBy('col')
df = df.withColumn("group", id().over(w))

有什么办法可以达到这样的目的吗？（我不能简单地使用 col 作为组 ID，因为我有兴趣在多个列上创建 window）

Answer 1

您可以使用原始数据帧为不同的 col 和自身 join 分配 row_number。

val data = Seq(
  (1, "a"),
  (2, "a"),
  (3, "b"),
  (4, "c"),
  (5, "c")
).toDF("id","col")

val df2 = data.select("col").distinct()
  .withColumn("group", row_number().over(Window.orderBy("col")))


val result = data.join(df2, Seq("col"), "left")
    .drop("col")

代码在 scala 中，但可以轻松更改为 pyspark。

希望这对您有所帮助

Answer 2

只需在 Window 函数上使用 dense_rank 内置函数即可得到您想要的结果

from pyspark.sql import window as W import pyspark.sql.functions as f df.select('id', f.dense_rank().over(W.Window.orderBy('col')).alias('group')).show(truncate=False)

哪个应该给你

+---+-----+ |id |group| +---+-----+ |1 |1 | |2 |1 | |3 |2 | |4 |3 | |5 |3 | +---+-----+

在 Spark Dataframe 的 window 上创建组 ID

Create a group id over a window in Spark Dataframe

window-functions

apache-spark

apache-spark-sql

pyspark