PySpark读取大量小Parquet文件性能优化：深入理解与解决方案

import pyspark
from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, StringType, IntegerType # 示例schema类型

# 初始化SparkSession
conf = pyspark.SparkConf().set('spark.driver.memory', '3g')
spark = (
    SparkSession.builder
    .master("local[10]") # 在本地模式下使用10个线程
    .config(conf=conf)
    .appName("Spark Local")
    .getOrCreate()
)

# 示例：假设已知Schema，或者从单个文件推断
# 实际场景中，如果所有文件Schema一致，可提前定义或从一个文件推断
# 例如：
# schema = StructType([
#     StructField("column1", StringType(), True),
#     StructField("column2", IntegerType(), True)
# ])
# 或者像问题中那样从一个文件推断：
df_sample = spark.read.parquet(r"C:\Project Data\Data-0.parquet")
schema = df_sample.schema
print("Schema successfully inferred from sample file.")
df_sample.printSchema()

# 尝试读取所有文件
# 假设文件路径模式为 "C:\Project Data\Data-*.parquet"
print("Attempting to read all partitioned parquet files using specified schema...")
df = spark.read.format("parquet") \
     .schema(schema) \
     .load(r"C:\Project Data\Data-*.parquet")

# 此时，即使没有立即触发Action，用户也可能观察到长时间的等待和内存消耗增加
# 例如，尝试执行一个Action：
# print(f"Total records: {df.count()}") # 这将触发实际计算
# df.show(5) # 或者显示前几行

# 假设df_raw是您通过上述慢速方式读取的DataFrame
# 这一步仍然会慢，但它将作为一次性的数据加载和转换过程
df_raw = spark.read.format("parquet") \
             .schema(schema) \
             .load(r"C:\Project Data\Data-*.parquet")

print(f"Successfully loaded initial DataFrame from small files.")
# df_raw.count() # 可以选择在这里触发count来获取总记录数

# 确定目标分区数
# 假设原始数据总大小为 1300 * 8MB = 10400MB (约10.4GB)
# 目标文件大小为 128MB/文件，则所需分区数约为 10400MB / 128MB = 81.25
# 我们可以选择一个合适的整数，例如 80 或 100
target_partitions = 80 # 根据总数据量和期望的文件大小进行调整

# 对DataFrame进行重分区，并将结果写入新的Parquet目录
# 这将生成大约 target_partitions 个较大的Parquet文件
print(f"Repartitioning data into {target_partitions} files and writing to new location...")
output_path = r"C:\Project Data\Consolidated_Data" # 新的存储路径

df_raw.repartition(target_partitions) \
      .write \
      .mode("overwrite")