如何在PySpark中根据出生日期计算星座（适配海量数据集）

需求说明 #

现有包含name（姓名）和dob（出生日期）字段的数据集，需基于给定的星座映射规则新增sign（星座缩写）列，因数据规模庞大，需使用PySpark实现分布式处理。

星座映射规则 #

zodiacs = [(120, 'Cap'), (218, 'Aqu'), (320, 'Pis'), (420, 'Ari'), (521, 'Tau'),
           (621, 'Gem'), (722, 'Can'), (823, 'Leo'), (923, 'Vir'), (1023, 'Lib'),
           (1122, 'Sco'), (1222, 'Sag'), (1231, 'Cap')]

映射逻辑：将出生日期转换为MMDD格式的整数，匹配对应区间的星座缩写（例如1114属于1023 < 1114 <= 1122，对应Sco）。

实现步骤 #

1. 初始化Spark环境并创建示例数据集

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, when, concat, lpad, month, dayofmonth

# 初始化SparkSession
spark = SparkSession.builder.appName("ZodiacSignProcessing").getOrCreate()

# 示例数据集
sample_data = [("John", "1932-11-14"), ("Maike", "1932-10-14")]
df = spark.createDataFrame(sample_data, ["name", "dob"])

2. 提取出生日期的月日数值
   将dob字段转换为日期类型后，提取月份和日期，拼接成MMDD格式的整数（例如11月14日转为1114）：

df = df.withColumn(
    "mmdd",
    concat(
        lpad(month(col("dob")), 2, "0"),  # 月份补前导零，确保两位
        lpad(dayofmonth(col("dob")), 2, "0")  # 日期补前导零，确保两位
    ).cast("int")
)

3. 构建星座匹配条件并新增列
   遍历星座映射列表，用PySpark内置的when函数构建条件链（避免使用UDF，提升大规模数据处理效率）：

# 定义星座映射列表
zodiacs = [(120, 'Cap'), (218, 'Aqu'), (320, 'Pis'), (420, 'Ari'), (521, 'Tau'),
           (621, 'Gem'), (722, 'Can'), (823, 'Leo'), (923, 'Vir'), (1023, 'Lib'),
           (1122, 'Sco'), (1222, 'Sag'), (1231, 'Cap')]

# 初始化第一个条件
sign_condition = when(col("mmdd") <= zodiacs[0][0], zodiacs[0][1])

# 遍历剩余映射规则，拼接条件
for idx in range(1, len(zodiacs)):
    prev_date, prev_sign = zodiacs[idx-1]
    curr_date, curr_sign = zodiacs[idx]
    sign_condition = sign_condition.when(
        (col("mmdd") > prev_date) & (col("mmdd") <= curr_date),
        curr_sign
    )

# 新增sign列，删除中间临时列mmdd
result_df = df.withColumn("sign", sign_condition).drop("mmdd")

4. 查看结果

result_df.show()

输出结果：

+-----+----------+----+
| name|       dob|sign|
+-----+----------+----+
| John|1932-11-14| Sco|
|Maike|1932-10-14| Lib|
+-----+----------+----+

注意事项 #

• 使用PySpark内置函数而非自定义UDF：内置函数经过分布式优化，在处理大规模数据集时性能更优，避免UDF带来的序列化/反序列化开销。
• 确保dob字段为日期类型：若原始数据中dob是字符串，可先通过to_date(col("dob"), "yyyy-MM-dd")转换为日期类型后再提取月日。

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