每完成 4 位数重置 pyspark 行号答案

【问题标题】：pyspark row number reset for every 4 digits completion每完成 4 位数重置 pyspark 行号
【发布时间】：2021-12-08 17:36:43
【问题描述】：

我基于一个列分区在特定列上应用 row_number。通过pyspark达到999后如何将行号重置为起始号1？

partcode	row number
12345	0001
12345	0002
12345	0003
12345	9999 -- Reaches max
12345	0001
12345	0002

【问题讨论】：

标签： pyspark row-number

【解决方案1】：

生成row number后：

从row number 中减去1 得到从零开始的行号
使用cut-off value (9999) 应用模运算
将1 添加到模结果以得到row number 从 1 开始并回绕到 1 的行号 cut-off value + 1

在我的示例中，我使用了较小的截止值3

from pyspark.sql import Window
from pyspark.sql import functions as F

df = spark.createDataFrame([(12345, )for _ in range(0, 10)], ("partcode", ))

window_spec = Window.partitionBy("partcode").orderBy("partcode")

cut_off = 3

df.withColumn("row number", F.row_number().over(window_spec))\
  .withColumn("row number", F.lpad(((F.col("row number") - F.lit(1)) % F.lit(cut_off)) + F.lit(1), 4, '0'))\
  .show()

输出

+--------+----------+
|partcode|row number|
+--------+----------+
|   12345|      0001|
|   12345|      0002|
|   12345|      0003|
|   12345|      0001|
|   12345|      0002|
|   12345|      0003|
|   12345|      0001|
|   12345|      0002|
|   12345|      0003|
|   12345|      0001|
+--------+----------+

【讨论】：