字符串列映射为整数时分区变为1的问题及替代方案咨询

一、替代row_number()的高效映射方案 #

1. 用Spark ML的StringIndexer（推荐）

这是Spark官方提供的字符串转整数索引工具，无需手动写窗口函数，还能自动处理重复值，关键是不会触发全局单分区的问题：

from pyspark.ml.feature import StringIndexer

# 初始化索引器，指定输入输出列，handleInvalid处理未见过的新值
indexer = StringIndexer(inputCol="st_id", outputCol="st_id_int", handleInvalid="keep")
# 拟合原数据生成映射规则，再转换
indexed_df = indexer.fit(your_dataframe).transform(your_dataframe)

默认它会按字符串出现频率分配索引（出现最多的为0），如果需要按字典序分配，可以先提取去重的st_id排序后再生成映射表。

2. 预生成映射表+Join（可控性强）

如果需要完全自定义映射规则（比如严格按字典序从1开始），可以先生成独立的映射表，再和原表关联，全程保持并行分区：

# 1. 提取去重的st_id，先重分区避免后续单节点压力
distinct_st = your_dataframe.select("st_id").distinct().repartition(16)  # 分区数按集群资源调整
# 2. 按字典序排序去重后的st_id
sorted_st = distinct_st.orderBy("st_id")
# 3. 用zipWithIndex生成索引，索引从0开始，要从1开始的话加1
st_mapping = sorted_st.rdd.zipWithIndex().toDF(["st_row", "st_id_int"])
st_mapping = st_mapping.select("st_row.st_id", (st_mapping["st_id_int"] + 1).alias("st_id_int"))
# 4. 和原表关联，保留原表分区结构
result_df = your_dataframe.join(st_mapping, on="st_id", how="left")

这个方法的优势是映射规则完全可控，且所有步骤都能并行执行，不会出现单分区瓶颈。

二、为什么row_number()会导致单分区？ #

row_number() over (order by st_id)是全局窗口函数，没有指定partition by子句，Spark会触发全局排序（将所有数据shuffle到一个节点）来生成连续的行号，这就是分区变为1的核心原因。后续的repartition只是把单分区的数据重新拆分，但前面的全局排序已经造成了巨大的性能开销，属于治标不治本。

三、如果一定要用窗口函数的补救方法 #

如果坚持要用窗口函数，必须给窗口加上partition by来避免全局排序，比如按st_id的哈希值分区，然后在每个分区内生成局部id，再通过全局偏移量拼接成唯一id，但这个方法比较繁琐，不如上面的方案高效：

from pyspark.sql.functions import hash, row_number, sum
from pyspark.sql.window import Window

# 1. 先按哈希值分区，避免全局排序
window_part = Window.partitionBy(hash("st_id") % 16).orderBy("st_id")
df_with_local_id = your_dataframe.withColumn("local_id", row_number().over(window_part))
# 2. 计算每个分区的id偏移量
window_global = Window.orderBy(hash("st_id") % 16)
df_with_offset = df_with_local_id.withColumn("offset", sum("local_id").over(window_global.rowsBetween(Window.unboundedPreceding, Window.currentRow - 1)))
# 3. 生成全局唯一id
result_df = df_with_offset.withColumn("st_id_int", df_with_offset["offset"] + df_with_offset["local_id"])

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