如何为DataFrame添加索引列并按指定列分组生成组内序号？

问题描述 #

原始DataFrame：

+-----+
|  Cur|
+-----+
|  HKD|
|  HKD|
|  HKD|
|  USD|
|  USD|
|  USD|
|  USD|
|  INR|
|  INR|
|  INR|
+-----+

期望生成包含分组内连续索引的结果：

+-----+-----+
|  Idx|  Cur|
+-----+-----+
|    1|  HKD|
|    2|  HKD|
|    3|  HKD|
|    1|  USD|
|    2|  USD|
|    3|  USD|
|    4|  USD|
|    1|  INR|
|    2|  INR|
|    3|  INR|
+-----+-----+

尝试过.sort()、.sortlevel()、.orderBy()等方法无法实现，需要按分组生成连续序号的方案。

方案1：PySpark 实现 #

利用窗口函数row_number()，按Cur分组后生成从1开始的连续序号：

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

# 假设df为原始DataFrame
window = Window.partitionBy("Cur").orderBy(F.lit(1))  # 用F.lit(1)保留原分组内的行顺序
result_df = df.withColumn("Idx", F.row_number().over(window))

# 输出结果
result_df.show()

如果不需要保留原分组内的行顺序，可以将orderBy(F.lit(1))替换为实际需要排序的字段，比如orderBy("Cur")。

方案2：Pandas 实现 #

使用groupby()结合cumcount()方法，cumcount()默认从0开始计数，加1后得到从1开始的序号：

import pandas as pd

# 假设df为原始DataFrame
df["Idx"] = df.groupby("Cur").cumcount() + 1

# 输出结果
print(df)

内容的提问来源于stack exchange，提问作者jasondesu