如何用Spark Structured Streaming监控S3 Raw顶层目录并分表写入？

1. 解决Parquet Schema无法推断的问题 #

监控raw/*这类多目录时，Spark无法自动统一推断所有Parquet文件的Schema（尤其是不同子目录Schema存在差异时），必须手动指定，两种可行方式：

方式一：预定义Schema #

若已知所有Parquet文件的结构，直接用StructType定义：

from pyspark.sql.types import StructType, StructField, StringType, IntegerType

# 示例Schema，根据实际数据结构调整
custom_schema = StructType([
    StructField("col1", IntegerType(), nullable=True),
    StructField("col2", StringType(), nullable=True),
    # 补充其他字段...
])

方式二：从样本文件加载Schema #

若不确定Schema，先读取一个现有Parquet文件获取结构：

# 读取任意存在的Parquet文件提取Schema
sample_df = self.spark.read.parquet("s3a://storage-layer/raw/GTest/abc/*")
custom_schema = sample_df.schema

2. 提取文件路径并映射到输出目录 #

用Spark内置函数input_file_name()获取每行数据对应的原始文件路径，再通过字符串替换将raw替换为processed，生成目标输出路径：

from pyspark.sql.functions import input_file_name, regexp_replace

stream_df = (
    self.spark
    .readStream
    .format("parquet")
    .schema(custom_schema)  # 手动指定Schema
    .option("maxFilesPerTrigger", 20)
    .load("s3a://storage-layer/raw/*/*/*")  # 通配符数量匹配实际目录深度，按需调整
    .withColumn("input_path", input_file_name())
    # 替换路径中的raw为processed，得到完整输出目录
    .withColumn("output_parent_dir", regexp_replace("input_path", "s3a://storage-layer/raw/(.*)/.*", "s3a://storage-layer/processed/$1"))
)

3. 使用foreachBatch实现分目录写入 #

在foreachBatch函数中，按批次提取唯一输出目录，将对应数据写入目标路径：

def process_batch(df, epoch_id):
    # 获取当前批次所有需写入的唯一目录
    output_dirs = df.select("output_parent_dir").distinct().collect()
    
    for row in output_dirs:
        target_dir = row["output_parent_dir"]
        # 筛选对应目录的数据，移除路径字段后写入
        df.filter(df.output_parent_dir == target_dir)
          .drop("input_path", "output_parent_dir")
          .write
          .mode("append")
          .parquet(target_dir)

# 启动流处理任务
(
    stream_df
    .writeStream
    .foreachBatch(process_batch)
    .outputMode("append")
    .trigger(processingTime="5 seconds")
    .start()
    .awaitTermination()
)

关键注意事项 #

• 目录通配符：load方法的通配符数量需匹配raw下的实际目录层级，若子目录更深，需增加通配符（如raw/*/*/*/*）。
• Schema一致性：若不同子目录的Parquet Schema不一致，需提前做Schema合并处理（流处理中mergeSchema选项支持有限，建议尽量保证所有目录Schema统一）。
• 写入模式：示例用append模式追加数据，可根据业务需求调整为overwrite等模式。

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