如何在PySpark中将2000万条数据集等分为2万条每组的命名分组？

为什么Bucketizer不适用 #

Bucketizer是针对连续数值列的范围分桶（比如将年龄划分为多个区间），而你的需求是按行数均分，两者的适用场景完全不同，所以用Bucketizer无法实现你的目标。

可行实现方案 #

我们可以通过添加全局行号（或唯一ID），再计算分组ID的方式来实现均分，具体步骤如下：

方案1：基于全局行号的严格均分（保证每组精确2万条）

这种方式会生成严格的每组2万条数据，适合对分组数量和大小有精确要求的场景：

from snowflake.snowpark import Window
from snowflake.snowpark.functions import row_number, floor, lit, concat

# 定义每组的大小
batch_size = 20000

# 1. 添加全局行号（用常量列排序保证全局唯一行号，不依赖原数据顺序）
window_spec = Window.order_by(lit(1))
df_with_row_num = df.with_column("ROW_NUM", row_number().over(window_spec))

# 2. 生成分组名称（Group1、Group2...）
df_with_group = df_with_row_num.with_column(
    "GROUP_NAME",
    concat(
        lit("Group"),
        # 计算分组ID：(行号-1)//每组大小 +1，确保从Group1开始
        floor((df_with_row_num["ROW_NUM"] - 1) / batch_size) + 1
    )
)

方案2：基于唯一ID的高效均分（近似2万条，性能更优）

如果不需要严格精确的每组2万条，只是希望拆分数据减少单次写入压力，可以用monotonically_increasing_id()生成分布式唯一ID，避免全局排序的性能开销：

from snowflake.snowpark.functions import monotonically_increasing_id, floor, concat, lit

batch_size = 20000

# 1. 添加分布式唯一ID
df_with_id = df.with_column("UNIQUE_ID", monotonically_increasing_id())

# 2. 生成分组名称
df_with_group = df_with_id.with_column(
    "GROUP_NAME",
    concat(
        lit("Group"),
        floor(df_with_id["UNIQUE_ID"] / batch_size) + 1
    )
)

遍历分组写入数据 #

生成分组后，就可以遍历每个分组执行写入操作，避免单次处理数据量过大超时：

# 获取所有分组名称
group_names = [row[0] for row in df_with_group.select("GROUP_NAME").distinct().collect()]

# 遍历每个分组写入
for group_name in group_names:
    # 筛选当前分组的数据
    group_df = df_with_group.filter(df_with_group["GROUP_NAME"] == group_name)
    
    # 执行你的写入逻辑，比如写入Snowflake表
    group_df.write.mode("append").save_as_table(f"YOUR_TARGET_TABLE_{group_name}")
    
    # 或者根据需求执行其他操作，比如导出文件等

注意事项 #

• 如果原数据已经有唯一排序键，可以用该列代替lit(1)进行排序，减少窗口函数的计算开销。
• 执行大表操作时，确保你的Snowflake仓库有足够的计算资源，避免任务排队或超时。

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