PySpark实现按用户循环创建临时表并生成合并抽样表

完全可以实现你需求的操作，下面是具体的实现步骤和代码示例：

一、按你的需求（循环创建临时表再合并）实现 #

1. 准备数据与获取用户样本量列表 #

首先读取存储用户任务的大表（假设名为user_tasks）和Table1，再将Table1中的用户ID与对应样本量转换为本地可迭代的列表：

# 读取大表和Table1
user_tasks_df = spark.table("user_tasks")
table1_df = spark.table("Table1")

# 将Table1转换为本地列表（适合用户数量不极大的场景）
user_sample_list = table1_df.select("user_id", "sample_size").collect()

2. 循环创建用户临时表 #

遍历每个用户，从大表中抽取对应数量的随机样本，并创建临时视图：

temp_table_names = []

for user in user_sample_list:
    user_id = user["user_id"]
    sample_size = user["sample_size"]
    
    # 抽取当前用户的随机样本
    user_sample_df = user_tasks_df.filter(f"user_id = '{user_id}'") \
                                 .orderBy(spark.rand()) \
                                 .limit(sample_size)
    
    # 创建临时表，命名格式示例：user_sample_{user_id}
    temp_table_name = f"user_sample_{user_id}"
    user_sample_df.createOrReplaceTempView(temp_table_name)
    temp_table_names.append(temp_table_name)

3. 合并所有临时表为总抽样表 #

通过UNION ALL语法合并所有临时表的数据，保存为最终结果表：

# 构建合并SQL语句
union_sql = " UNION ALL ".join([f"SELECT * FROM {table}" for table in temp_table_names])
final_sample_df = spark.sql(union_sql)

# 将结果保存为永久表（也可保存为临时表，根据需求调整）
final_sample_df.write.mode("overwrite").saveAsTable("final_user_task_samples")

4. 删除临时表 #

遍历临时表名称列表，删除所有创建的临时视图：

for table_name in temp_table_names:
    spark.catalog.dropTempView(table_name)

二、更高效的分布式实现方案（推荐） #

上述循环方式会触发多次Spark作业，当用户数量较多时效率偏低。更贴合Spark分布式特性的做法是通过窗口函数结合Table1的样本量，一次性完成抽样：

# 关联大表和Table1，获取每个用户需要抽取的样本量
joined_df = user_tasks_df.join(table1_df, on="user_id", how="inner")

# 对每个用户的任务随机排序，取前N个样本
from pyspark.sql.window import Window
from pyspark.sql.functions import row_number, rand

window_spec = Window.partitionBy("user_id").orderBy(rand())

final_sample_df = joined_df.withColumn("row_num", row_number().over(window_spec)) \
                          .filter("row_num <= sample_size") \
                          .drop("row_num", "sample_size")

# 保存结果
final_sample_df.write.mode("overwrite").saveAsTable("final_user_task_samples")

这种方式仅需一次Spark作业，性能远优于循环处理，建议优先使用。

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