Pyspark：如何添加带有行号的列？

Question

我有一个pyspark数据框。我想添加一个包含行号的列。

这就是我正在做的

stop_df = stop_df.withColumn("stop_id", monotonically_increasing_id())

如果我检查 的最大值stop_id，我得到

stop_df.agg(max("stop_id")).show()
+--------------+
|  max(stop_id)|
+--------------+
|32478542692458|
+--------------+

而行数是

stop_df.count()
Out[4]: 8134605

Answer 1

来自 spark monotonically_increasing_id文档：

生成单调递增的 64 位整数rs 的列。

生成的 ID 保证单调递增且唯一，但不连续。当前实现将分区 ID 放在高 31 位，将每个分区内的记录号放在低 33 位。假设数据帧的分区数少于 10 亿，每个分区的记录数少于 80 亿。

使用窗口 row_number函数获取行号。

df=spark.createDataFrame([("a",),("b",)],["id"])
from pyspark.sql.window import Window
from pyspark.sql.functions import *
#add partition by and order by clause if ordering required with in window.
w=Window.orderBy(lit(1))

df.withColumn("stop_id",row_number().over(w)).show()
#+---+-------+
#| id|stop_id|
#+---+-------+
#|  a|      1|
#|  b|      2|
#+---+-------+

df.withColumn("stop_id",row_number().over(w)).agg(max("stop_id")).show()
#+------------+
#|max(stop_id)|
#+------------+
#|           2|
#+------------+

df.count()
#2