基于年度日期区间计算列汇总值：Spark RDD与DataFrame实现需求

问题描述 #

现有DataFrame df1：

需要计算2021、2022、2023各年度的产品平均价格，期望输出：

计算逻辑 #

• 产品P1 2021年：((6/7/2021 - 1/1/2021)12 + (10/19/2021 - 6/8/2021)14 + (12/31/2021 - 10/20/2021)13)/365 = (15812 + 13314 + 5313)/365 = 12.18
• 产品P1 2022年：((5/14/2022 - 1/1/2022)13 + (11/20/2022 - 5/15/2022)12.5 + (12/31/2022 - 11/21/2022)12.5)/365 = (13413 + 18912.5 + 4012.5)/365 = 12.61
• 产品P1 2023年：12.5（自2022年11月21日起价格未变，全年保持该值）

尝试用RDD的map函数处理但未得到正确结果，寻求解决方案。

解决方案 #

使用PySpark DataFrame API实现，相比RDD更简洁高效，核心思路是计算价格区间与目标年份的重叠天数，再通过加权平均得到年度价格：

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.types import DateType

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

# 构造示例数据
data = [
    ("p1", "6/12/2020", "6/7/2021", 12),
    ("p1", "6/8/2021", "10/19/2021", 14),
    ("p1", "10/20/2021", "5/14/2022", 13),
    ("p1", "5/15/2022", "11/20/2022", 12.5),
    ("p1", "11/21/2022", "1/1/2099", 12.5),
    ("p2", "6/12/2017", "6/7/2021", 12),
    ("p2", "6/8/2021", "10/19/2021", 14),
    ("p2", "10/20/2021", "5/14/2022", 12),
    ("p2", "5/15/2022", "5/14/2099", 11)
]

df1 = spark.createDataFrame(data, ["产品", "开始日期", "结束日期", "价格"])

# 转换日期列为Spark日期类型
df1 = df1.withColumn("开始日期", F.to_date("开始日期", "M/d/yyyy")) \
         .withColumn("结束日期", F.to_date("结束日期", "M/d/yyyy"))

# 生成目标年份的起止日期
target_years = [2021, 2022, 2023]
year_dates = spark.createDataFrame(
    [(y, f"{y}-01-01", f"{y}-12-31") for y in target_years], 
    ["年份", "年度开始", "年度结束"]
)
year_dates = year_dates.withColumn("年度开始", F.to_date("年度开始")) \
                       .withColumn("年度结束", F.to_date("年度结束"))

# 交叉连接，关联每个价格区间与目标年份
cross_df = df1.crossJoin(year_dates)

# 计算价格区间与年度的重叠起止日期
cross_df = cross_df.withColumn("重叠开始", F.greatest("开始日期", "年度开始")) \
                   .withColumn("重叠结束", F.least("结束日期", "年度结束"))

# 筛选出有有效重叠的记录
valid_df = cross_df.filter(F.col("重叠开始") <= F.col("重叠结束"))

# 计算重叠天数（包含首尾日期）
valid_df = valid_df.withColumn("重叠天数", F.datediff("重叠结束", "重叠开始") + 1)

# 计算价格加权值（价格 * 重叠天数）
valid_df = valid_df.withColumn("价格加权", F.col("价格") * F.col("重叠天数"))

# 按产品和年份分组，计算加权平均价格
result_df = valid_df.groupBy("产品", "年份") \
                    .agg(
                        F.round(F.sum("价格加权") / F.sum("重叠天数"), 2).alias("平均价格")
                    )

# 补充全年无分段的年度记录（如2023年）
# 获取每个产品的最新有效价格
latest_price = df1.groupBy("产品") \
                  .agg(F.max("开始日期").alias("最新日期")) \
                  .join(df1, ["产品", "开始日期"], "inner") \
                  .select("产品", "价格")

# 生成缺失的年度记录
missing_years = year_dates.join(result_df, ["年份"], "left_anti") \
                          .crossJoin(latest_price) \
                          .select("产品", "年份", F.col("价格").alias("平均价格"))

# 合并最终结果并排序
final_result = result_df.union(missing_years).orderBy("产品", "年份")

# 展示结果
final_result.show()

代码说明 #

1. 日期转换：将字符串日期转为Spark可计算的日期类型，避免字符串操作的误差。
2. 交叉连接：把每个价格区间与目标年份关联，找出区间在对应年份的重叠部分。
3. 重叠计算：用greatest和least确定重叠区间的起止，计算实际有效天数。
4. 加权平均：通过价格乘以天数得到加权值，求和后除以总天数，得到年度平均价格。
5. 补充缺失记录：对于全年处于同一价格区间的年份，直接取产品最新价格作为年度平均。

运行代码后将得到与期望一致的输出结果。

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