如何为Spark DataFrame中的企业分配唯一编号？

我有一个包含Company（企业）和contract number（合同编号）字段的数据集，想要新增一个名为Company number的列，为每个企业分配唯一编号，预期效果如下：

Company | contract number | Company number
amazon  | IO123456        | 1
google  | IO113456        | 2
google  | IO111456        | 2
yahoo   | IO111156        | 3
amazon  | IO111116        | 1

我尝试了以下代码：

window_spec = Window.partitionBy('Company')
Programmatic_df = Programmatic_df.withColumn('Company number', dense_rank().over(window_spec))

但没有得到预期结果，这段代码给每个企业的每行分配了索引而非唯一编号，实际结果如下：

Company | contract number | Company number
amazon  | IO123456        | 1
google  | IO113456        | 1
google  | IO111456        | 2
yahoo   | IO111156        | 1
amazon  | IO111116        | 2

请问正确的实现方法是什么？

你之前的代码错误在于用partitionBy('Company')把窗口按企业拆分，dense_rank()会在每个企业内部做行级排名，因此同一企业的多行编号会递增。要实现每个企业对应唯一编号，有两种可行方案：

方案1：全局窗口+dense_rank() #

定义不分区的全局窗口，按Company排序后使用dense_rank()，这样每个唯一企业会获得一个全局唯一的编号，同一企业的所有行编号一致：

from pyspark.sql.window import Window
from pyspark.sql.functions import dense_rank

# 全局窗口按Company排序
window_spec = Window.orderBy('Company')
Programmatic_df = Programmatic_df.withColumn('Company number', dense_rank().over(window_spec))

编号会按照Company的字典序递增，完全符合你的预期效果。

方案2：提取唯一企业生成编号后关联 #

如果需要自定义编号顺序（比如不依赖字典序），可以先提取唯一企业列表生成编号，再关联回原数据集：

from pyspark.sql.functions import monotonically_increasing_id

# 提取唯一企业并生成从1开始的编号
unique_companies = Programmatic_df.select('Company').distinct()\
    .orderBy('Company')\
    .withColumn('Company number', monotonically_increasing_id() + 1)

# 关联回原数据集，保留所有行
Programmatic_df = Programmatic_df.join(unique_companies, on='Company', how='left')

这里用monotonically_increasing_id()生成递增ID，加1是为了让编号从1开始，你也可以根据需求替换成自定义的编号规则。

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