如何实现DataFrame条件分区，确保同一col1值的行处于同一分区？

我有一个包含两列的DataFrame，数据如下：

col1 | col2
------------
a1   |   b1
------------
a2   |   b1
------------
a3   |   b2
------------
a1   |   b2
------------
a1   |   b3
------------

我原本通过生成随机数的方式对该DataFrame进行分区，代码如下：

df = df.withColumn("part", (rand() * num_partitions).cast("int"))
df.write.partitionBy("part").mode("overwrite").parquet("/address/")

但这种分区方式无法保证所有col1=a1的行都被分配到同一个分区，请问有没有方法能在分区时实现这一保证？

当然可以实现，核心是基于col1的哈希值计算分区ID——相同col1值的行哈希结果一致，取模后必然得到同一个分区编号，就能保证它们被分到同一分区。

具体实现代码 #

from pyspark.sql.functions import hash, abs

# 假设预设分区数为num_partitions
df = df.withColumn("part", abs(hash("col1")) % num_partitions)
df.write.partitionBy("part").mode("overwrite").parquet("/address/")

关键说明 #

• hash("col1")会生成对应列值的哈希码，由于哈希码可能为负，用abs()取绝对值避免出现负分区号
• 对哈希值取模num_partitions，能把分区号限制在0到num_partitions-1的范围内，匹配你需要的分区数量
• 如果col1的不同取值数量远大于分区数，会出现多个不同col1值共享同一分区的情况，但这是分区数有限时的正常现象，不会破坏"相同col1值在同一分区"的核心要求

内容的提问来源于stack exchange，提问作者A.M.