PySpark如何输出按指定字段排序的CSV/Parquet文件？

需求：从大型BigQuery表读取数据，导出为多个CSV文件，每个文件包含61000条全局有序的记录（按colA、colB排序）。

当前尝试的代码片段：

TMP_BUCKET = "stg-gcs-bucket"
MAX_PARTITION_BYTES = str(512 * 1024 * 1024)
MAX_ROW_NUM_PER_FILE = "61000"

spark = SparkSession \
    .builder \
    .master('yarn') \
    .appName('crs-bq-export-csv') \
    .config('spark.sql.execution.arrow.pyspark.enabled', 'true') \
    .config('spark.jars', 'gs://spark-lib/bigquery/spark-bigquery-with-dependencies_2.12-0.23.2.jar') \
    .config("spark.sql.broadcastTimeout", "36000") \
    .config("spark.sql.files.maxRecordsPerFile", MAX_ROW_NUM_PER_FILE) \
    .config("spark.sql.files.maxPartitionBytes", MAX_PARTITION_BYTES) \
    .config("spark.files.maxPartitionBytes", MAX_PARTITION_BYTES) \
    .config("spark.driver.maxResultSize", "24g") \
    .config("spark.sql.execution.arrow.pyspark.enabled", "true") \
    .getOrCreate()

# 读取BigQuery数据
df = spark.read.format('bigquery') \
    .option('table', TABLE_NAME) \
    .load()

# 尝试排序后写入
df.sort('colA').sort('colB').write.mode('overwrite').csv(OUTPUT_PATH, header=True)

遇到的问题：生成的CSV文件记录杂乱无章，未按colA、colB全局排序：

• 预期的有序格式：

colA colB
1 1
2 2
3 3
...
60001 60001

• 实际得到的无序格式：

colA colB
2 1
3 3
2 2
1 3

原因分析 #

1. 多次排序覆盖问题：df.sort('colA').sort('colB')只会保留最后一次排序规则（仅按colB排序），无法实现多列联合排序。
2. Spark排序的局限性：默认的sort/orderBy仅保证分区内有序，但Spark写入文件时会并行处理各个分区，分区之间的顺序不做保证，因此全局数据呈现无序状态。
3. shuffle操作的影响：Spark为提升性能会对数据进行shuffle，打乱全局排序结果，最终写入的文件顺序无法匹配全局排序逻辑。

解决方案：实现全局有序的分区导出 #

要生成全局有序的分区CSV，需要先完成全局排序，再基于排序结果进行精准分区，确保每个分区的记录数正好是61000条，且分区内保持有序。具体步骤如下：

步骤1：修正多列排序逻辑 #

使用orderBy(colA, colB)一次指定多列排序规则，避免多次排序覆盖的问题。

步骤2：添加全局递增行号 #

通过窗口函数为全局排序后的每条数据添加唯一递增的行号，用于后续计算分区分组。

步骤3：按行号计算分区键 #

用行号除以61000取整，得到每个记录所属的文件分组ID，确保每个分组正好包含61000条记录。

步骤4：按分区键重分区并写入 #

基于分组键进行重分区，确保每个分区对应一个输出文件；同时在分区内再次排序，保证文件内记录有序，最后写入CSV。

修改后的完整代码 #

from pyspark.sql.window import Window
from pyspark.sql.functions import row_number, floor, col

TMP_BUCKET = "stg-gcs-bucket"
MAX_PARTITION_BYTES = str(512 * 1024 * 1024)
MAX_ROW_NUM_PER_FILE = 61000  # 改为整数，避免字符串转换问题

spark = SparkSession \
    .builder \
    .master('yarn') \
    .appName('crs-bq-export-csv') \
    .config('spark.sql.execution.arrow.pyspark.enabled', 'true') \
    .config('spark.jars', 'gs://spark-lib/bigquery/spark-bigquery-with-dependencies_2.12-0.23.2.jar') \
    .config("spark.sql.broadcastTimeout", "36000") \
    .config("spark.driver.maxResultSize", "24g") \
    .getOrCreate()

# 读取BigQuery数据
df = spark.read.format('bigquery') \
    .option('table', TABLE_NAME) \
    .load()

# 1. 全局排序并添加递增行号
window_spec = Window.orderBy(col('colA'), col('colB'))
df_with_rowid = df.withColumn('row_id', row_number().over(window_spec))

# 2. 计算每个记录所属的文件分组ID
df_with_fileid = df_with_rowid.withColumn(
    'file_id',
    floor((col('row_id') - 1) / MAX_ROW_NUM_PER_FILE)
)

# 3. 按file_id重分区，确保每个分区对应一个文件；分区内按row_id排序保证文件内有序
df_with_fileid.repartition(col('file_id')) \
    .sortWithinPartitions('row_id') \
    .drop('row_id', 'file_id')  # 移除临时列
    .write.mode('overwrite') \
    .csv(OUTPUT_PATH, header=True)

关键说明 #

• 全局排序+行号：通过窗口函数生成的row_id是全局连续递增的，确保数据的全局顺序。
• 分区键计算：floor((row_id -1)/61000)确保第1-61000条记录属于file_id=0，第61001-122000条属于file_id=1，以此类推，每个分组正好61000条记录。
• repartition+sortWithinPartitions：repartition(col('file_id'))让每个file_id对应一个Spark分区，写入时每个分区生成一个CSV文件；sortWithinPartitions('row_id')保证每个文件内的记录是有序的。
• 移除临时列：最后通过drop移除row_id和file_id，避免输出文件包含多余字段。

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