Pyspark数据帧按字母顺序拆分并写入S3

from pyspark.sql import SparkSession

# 创建SparkSession
spark = SparkSession.builder \
    .appName("Split and Write to S3") \
    .getOrCreate()

# 读取数据为DataFrame
data = spark.read.csv("input.csv", header=True, inferSchema=True)

# 按照指定列排序
sorted_data = data.sort("column_name")

# 获取排序后的列的首字母
first_letters = sorted_data.selectExpr("substring(column_name, 1, 1) as first_letter").distinct()

# 按照首字母拆分数据帧并写入S3
for letter_row in first_letters.collect():
    letter = letter_row.first_letter
    letter_data = sorted_data.filter(sorted_data.column_name.startswith(letter))
    letter_data.write.csv("s3://bucket_name/path/{}".format(letter))

# 关闭SparkSession
spark.stop()