Pyspark如何基于多列分组为相同位置生成唯一location_id新列

你可以通过以下两种常用方式实现需求，可根据数据量大小选择：

方案1：唯一位置维度表关联（推荐大数据量使用） #

该方案先提取所有不重复的位置组合生成唯一ID，再关联回原表，避免全表单分区shuffle，性能更稳定：

from pyspark.sql.functions import monotonically_increasing_id

# 1. 提取唯一位置组合，生成对应的location_id
location_dim = df.select("street", "locality", "town_city", "district", "country") \
    .distinct() \
    .withColumn("location_id", monotonically_increasing_id())

# 2. 关联原表，为每一行匹配对应位置的location_id
result_df = df.join(
    location_dim,
    on=["street", "locality", "town_city", "district", "country"],
    how="left"
)

# 可选：调整列顺序，将location_id放到最前
result_df = result_df.select(
    "location_id", "id", "price", "street", "locality", "town_city", "district", "country"
)

方案2：窗口函数实现（适合小数据量） #

如果数据量不大，可以用dense_rank窗口函数快速实现，代码更简洁：

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

# 定义窗口，按5个位置字段排序
window_spec = Window.orderBy("street", "locality", "town_city", "district", "country")

# 生成连续的location_id，减1是为了从0开始编号，和示例效果一致
result_df = df.withColumn("location_id", dense_rank().over(window_spec) - 1)

注意事项 #

• 方案1的monotonically_increasing_id生成的ID是全局唯一但不连续的，如果你需要连续的ID，可以把monotonically_increasing_id替换为窗口函数row_number().over(Window.orderBy("street", "locality", "town_city", "district", "country")) - 1
• 方案2在数据量超过百万级时容易出现单分区内存溢出，优先选择方案1

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