基于AWS EMR，如何用PySpark对S3分区Parquet按用户ID月度采样过滤

Hey there! 针对你要处理S3上大型分区Parquet数据集，在EMR上按user_id抽取10万用户样本并聚合的需求，我来分享一套高效的实现方案——重点是利用Spark的分区特性和优化策略，避免全量扫描，帮你节省资源和时间。

核心思路：先抽样本ID，再关联筛选数据 #

直接全量扫描每个分区去筛选会非常低效，更聪明的做法是先从全量数据中提取10万个唯一的user_id样本，再用这些ID去精准过滤每个分区的数据，大幅减少需要处理的数据量。

具体实现步骤 #

1. 初始化适配EMR的SparkSession #

先根据EMR集群的配置初始化SparkSession，调整关键参数保证大数据处理的稳定性：

from pyspark.sql import SparkSession
from pyspark.sql.functions import col

spark = SparkSession.builder \
    .appName("S3ParquetUserSample") \
    .config("spark.sql.shuffle.partitions", "200")  # 建议设为集群核心数的2-3倍
    .config("spark.driver.memory", "16g")  # 根据EMR主节点配置调整
    .config("spark.executor.memory", "32g")  # 根据EMR核心节点配置调整
    .getOrCreate()

2. 抽取10万个唯一的user_id样本 #

这里要避免直接对全量数据做distinct（会非常耗时），可以结合分区特性先采样再去重：

# 先读取数据集元数据，不加载实际数据
full_data_df = spark.read.parquet("s3://your-bucket/your-partitioned-parquet-path/")

# 方式1：如果user_id分布均匀，直接采样后取前10万
user_sample_df = full_data_df.select("user_id") \
    .distinct() \
    .sample(fraction=0.01, seed=42)  # 先按比例采样，再限制数量
    .limit(100000)

# 方式2：如果user_id分布不均，先从每个分区抽样本再合并去重
# user_sample_df = full_data_df.select("user_id") \
#     .sampleBy("month", fractions={"2024-01":0.005, "2024-02":0.005}, seed=42) \
#     .distinct() \
#     .limit(100000)

# 将样本ID转为广播变量，让每个Executor节点缓存这个列表，避免Shuffle传输
user_sample_broadcast = spark.sparkContext.broadcast(
    user_sample_df.rdd.map(lambda row: row.user_id).collect()
)

3. 读取分区数据并精准筛选样本用户 #

利用Spark的过滤下推特性，让Parquet存储层直接过滤掉不符合条件的行，避免全量加载：

# 用广播变量的ID列表做过滤，这个条件会自动下推到Parquet层面
filtered_df = spark.read.parquet("s3://your-bucket/your-partitioned-parquet-path/") \
    .filter(col("user_id").isin(user_sample_broadcast.value))

# 如果你需要按月份单独处理分区，也可以遍历分区路径读取，进一步优化：
# from pyspark.sql import DataFrame
# months = ["2024-01", "2024-02", "2024-03"]
# filtered_dfs = []
# for month in months:
#     month_df = spark.read.parquet(f"s3://your-bucket/your-path/month={month}/") \
#         .filter(col("user_id").isin(user_sample_broadcast.value))
#     filtered_dfs.append(month_df)
# filtered_df = spark.union(filtered_dfs)

4. 聚合计算并写回S3 #

根据你的业务需求完成聚合，再将结果写回S3（建议继续按分区存储，方便后续查询）：

# 示例：按user_id和月份统计行为次数
aggregated_df = filtered_df.groupBy("user_id", "month") \
    .count() \
    .withColumnRenamed("count", "total_activities")

# 写回S3，用Parquet格式+月份分区，覆盖已有数据
aggregated_df.write \
    .mode("overwrite") \
    .partitionBy("month") \
    .parquet("s3://your-bucket/your-output-path/")

关键优化Tips（针对超大型数据集） #

• 强制过滤下推：确保user_id是Parquet的列（不是嵌套字段），Spark会自动将过滤逻辑下推到存储层，只读取符合条件的行。
• 广播变量复用：把10万user_id做成广播变量，避免在Shuffle阶段重复传输这个列表，大幅减少网络开销。
• 分区并行处理：如果数据按月份分区，单独读取每个分区并过滤，能让集群的Executor节点并行处理不同分区，提升效率。
• 避免全量扫描：不要直接对全量数据做distinct，先采样再去重，能把计算量缩小几个数量级。

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