如何基于另一列在数据框中添加自动增量列

问题描述

我有一个类似于下面的 pyspark 数据框：

order_id    item    qty
123          abc    1
123          abc1   4
234          abc2   5
234          abc3   2
234          abc4   7
123          abc5   5
456          abc6   9
456          abc7   8
456          abc8   9

我想根据列“order_id”添加一个自动递增列，预期结果是：

order_id    item    qty AutoIncrementingColumn_orderID
123         abc      1   1
123         abc1     4   2
234         abc2     5   1
234         abc3     2   2
234         abc4     7   3
123         abc5     5   3
456         abc6     9   1
456         abc7     8   2
456         abc8     9   3

我找不到基于另一列生成的解决方案，知道如何实现吗？

解决方法

您可以使用row_number：

from pyspark.sql import functions as F,Window

df2 = df.withColumn(
    'AutoIncrementingColumn_orderID',F.row_number().over(Window.partitionBy('order_id').orderBy('item'))
)

df2.show()
+--------+----+---+------------------------------+
|order_id|item|qty|AutoIncrementingColumn_orderID|
+--------+----+---+------------------------------+
|     234|abc2|  5|                             1|
|     234|abc3|  2|                             2|
|     234|abc4|  7|                             3|
|     456|abc6|  9|                             1|
|     456|abc7|  8|                             2|
|     456|abc8|  9|                             3|
|     123| abc|  1|                             1|
|     123|abc1|  4|                             2|
|     123|abc5|  5|                             3|
+--------+----+---+------------------------------+

,

几种方法：

这里是sql方式：

df=Ss.sql("""
select order_id,item,qty,row_number() over(partition by order_id order by qty) as autoInc
from (
select order_id,qty
from ( values 
(123,'abc',1 ),(123,'abc1',4),(234,'abc2',5),'abc3',2),'abc4',7),'abc5',(456,'abc6',9),'abc7',8),'abc8',9)
) as T(order_id,qty))""")

df.show()

输出：

+--------+----+---+-------+
|order_id|item|qty|autoInc|
+--------+----+---+-------+
|     456|abc7|  8|      1|
|     456|abc6|  9|      2|
|     456|abc8|  9|      3|
|     234|abc3|  2|      1|
|     234|abc2|  5|      2|
|     234|abc4|  7|      3|
|     123| abc|  1|      1|
|     123|abc1|  4|      2|
|     123|abc5|  5|      3|
+--------+----+---+-------+