如何用Apache Spark或类似方案实现按日分摊求和？

需求概述 #

现有一批包含id、start_day、end_date、sum字段的记录，需将每条记录的sum值平均分摊到其日期区间内的每一天（区间为左闭右开，例如2022-12-01至2022-12-03仅包含1号、2号两天），输出格式为：
id, part_id, date, 每日分摊值

示例输入 #

1, 2022-12-01, 2022-12-03, 12
2, 2022-12-05, 2022-12-10, 100

示例输出 #

1, 1, 2022-12-01, 6
1, 2, 2022-12-02, 6
2, 1, 2022-12-05, 20
2, 2, 2022-12-06, 20
2, 3, 2022-12-07, 20
2, 4, 2022-12-08, 20
2, 5, 2022-12-09, 20

此前尝试RDBMS SQL方案性能不佳，不愿使用Java MapReduce实现，以下提供Apache Spark的高效分布式实现方案。

Apache Spark 实现（Scala版本） #

核心步骤 #

1. 加载数据并转换字段类型，确保日期、数值类型正确
2. 计算区间天数与每日分摊值
3. 生成日期序列并拆分为单行记录
4. 对每个id下的日期排序并生成part_id

代码示例 #

import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.{DateType, IntegerType, LongType}

// 加载输入数据（支持本地文件、HDFS、数据库等数据源）
val inputDF = spark.read
  .option("header", "false")
  .option("delimiter", ",")
  .csv("/path/to/input.csv")
  .toDF("id", "start_day", "end_date", "sum")
  // 字段类型转换
  .withColumn("id", col("id").cast(IntegerType))
  .withColumn("start_day", trim(col("start_day")).cast(DateType))
  .withColumn("end_date", trim(col("end_date")).cast(DateType))
  .withColumn("sum", col("sum").cast(LongType))

// 计算分摊值并生成日期序列
val explodedDF = inputDF
  // 计算区间天数（左闭右开，直接用datediff得到有效天数）
  .withColumn("days_count", datediff(col("end_date"), col("start_day")))
  // 过滤无效区间（天数为0或负数）
  .filter(col("days_count") > 0)
  // 计算每日分摊值
  .withColumn("daily_sum", col("sum") / col("days_count"))
  // 生成区间内的日期序列（从start_day到end_date前一天）
  .withColumn("date_sequence", sequence(col("start_day"), date_sub(col("end_date"), 1)))
  // 拆分日期序列为单行记录
  .select("id", "daily_sum", "date_sequence")
  .withColumn("date", explode(col("date_sequence")))
  .drop("date_sequence")

// 生成part_id并整理结果
val resultDF = explodedDF
  .groupBy("id")
  .agg(collect_list(struct("date", "daily_sum")).alias("date_list"))
  // 对日期列表按升序排序
  .withColumn("date_list", sort_array(col("date_list"), asc = true))
  // 按位置拆分并生成part_id（posexplode从0开始，需+1）
  .withColumn("part_date", explode(posexplode(col("date_list")).alias("part_id", "date_info")))
  .select(
    col("id"),
    (col("part_id") + 1).alias("part_id"),
    col("date_info.date").alias("date"),
    col("date_info.daily_sum").alias("daily_sum")
  )
  .orderBy("id", "part_id")

// 输出结果（可保存为CSV、Parquet或写入数据库）
resultDF.write
  .option("header", "false")
  .option("delimiter", ",")
  .csv("/path/to/output.csv")

Apache Spark 实现（Python版本） #

逻辑与Scala版本一致，使用PySpark实现：

代码示例 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import *
from pyspark.sql.types import IntegerType, DateType, LongType

# 初始化Spark会话
spark = SparkSession.builder.appName("DateSumSplit").getOrCreate()

# 加载输入数据
input_df = spark.read \
    .option("header", "false") \
    .option("delimiter", ",") \
    .csv("/path/to/input.csv") \
    .toDF("id", "start_day", "end_date", "sum") \
    .withColumn("id", col("id").cast(IntegerType)) \
    .withColumn("start_day", trim(col("start_day")).cast(DateType)) \
    .withColumn("end_date", trim(col("end_date")).cast(DateType)) \
    .withColumn("sum", col("sum").cast(LongType))

# 计算分摊值并生成日期序列
exploded_df = input_df \
    .withColumn("days_count", datediff(col("end_date"), col("start_day"))) \
    .filter(col("days_count") > 0) \
    .withColumn("daily_sum", col("sum") / col("days_count")) \
    .withColumn("date_sequence", sequence(col("start_day"), date_sub(col("end_date"), 1))) \
    .select("id", "daily_sum", "date_sequence") \
    .withColumn("date", explode(col("date_sequence"))) \
    .drop("date_sequence")

# 生成part_id并整理结果
result_df = exploded_df \
    .groupBy("id") \
    .agg(collect_list(struct("date", "daily_sum")).alias("date_list")) \
    .withColumn("date_list", sort_array(col("date_list"), asc=True)) \
    .withColumn("part_date", explode(posexplode(col("date_list")).alias("part_id", "date_info"))) \
    .select(
        col("id"),
        (col("part_id") + 1).alias("part_id"),
        col("date_info.date").alias("date"),
        col("date_info.daily_sum").alias("daily_sum")
    ) \
    .orderBy("id", "part_id")

# 输出结果
result_df.write \
    .option("header", "false") \
    .option("delimiter", ",") \
    .csv("/path/to/output.csv")

性能优化建议 #

• 分区调整：根据数据量设置spark.sql.shuffle.partitions参数（默认200），避免小文件或shuffle瓶颈
• 类型优化：确保字段类型正确，减少不必要的类型转换开销
• 提前过滤：过滤掉days_count <=0的无效记录，避免无效计算
• 存储格式：使用Parquet等列式存储格式替代CSV，提升读写与查询性能

小规模数据替代方案（Pandas） #

如果数据量较小，可使用Pandas快速实现：

import pandas as pd
from datetime import timedelta

# 加载数据
df = pd.read_csv("/path/to/input.csv", names=["id", "start_day", "end_date", "sum"])
df["start_day"] = pd.to_datetime(df["start_day"]).dt.date
df["end_date"] = pd.to_datetime(df["end_date"]).dt.date

# 生成分摊记录
result = []
for _, row in df.iterrows():
    days_count = (row["end_date"] - row["start_day"]).days
    if days_count <= 0:
        continue
    daily_sum = row["sum"] / days_count
    current_date = row["start_day"]
    for part_id in range(1, days_count + 1):
        result.append([row["id"], part_id, current_date, daily_sum])
        current_date += timedelta(days=1)

# 保存结果
result_df = pd.DataFrame(result, columns=["id", "part_id", "date", "daily_sum"])
result_df.to_csv("/path/to/output.csv", header=False, index=False)

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