如何将DataFrame中跨月起止日期拆分为每月分段行

需求示例 #

假设输入DataFrame结构如下：

需要转换为：

方案1：Pandas（适用于内存可容纳的数据） #

如果3000万条数据能塞进内存，用矢量化操作替代循环，避免性能瓶颈：

import pandas as pd

# 1. 确保日期列是datetime类型
df = pd.DataFrame({
    "ID": ["ID1", "ID2", "ID3"],
    "startDate": ["2023-02-15", "2023-05-20", "2023-12-25"],
    "endDate": ["2023-04-17", "2023-05-25", "2024-02-10"]
})
df["startDate"] = pd.to_datetime(df["startDate"])
df["endDate"] = pd.to_datetime(df["endDate"])

# 2. 生成每个ID对应的月度起始日期序列
df["month_starts"] = df.apply(
    lambda row: pd.date_range(
        start=row["startDate"].replace(day=1),
        end=row["endDate"],
        freq="MS"
    ),
    axis=1
)

# 3. 展开序列为多行
df_exploded = df.explode("month_starts")

# 4. 计算每个月度子区间的起止日期
df_exploded["new_start"] = df_exploded.apply(
    lambda row: max(row["startDate"], row["month_starts"]),
    axis=1
)
df_exploded["new_end"] = df_exploded.apply(
    lambda row: min(row["endDate"], row["month_starts"] + pd.offsets.MonthEnd(0)),
    axis=1
)

# 5. 整理最终结果
result = df_exploded.drop(columns=["startDate", "endDate", "month_starts"]).rename(
    columns={"new_start": "startDate", "new_end": "endDate"}
)
print(result)

方案2：PySpark（千万级数据首选，分布式处理） #

3000万条数据用Pandas容易出现内存溢出，推荐用PySpark分布式处理，性能更稳定：

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

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

# 构造示例数据（实际可从数据源读取）
data = [
    ("ID1", "2023-02-15", "2023-04-17"),
    ("ID2", "2023-05-20", "2023-05-25"),
    ("ID3", "2023-12-25", "2024-02-10")
]
df = spark.createDataFrame(data, ["ID", "startDate", "endDate"])

# 1. 转换日期列类型为DateType
df = df.withColumn("startDate", F.to_date("startDate"))
df = df.withColumn("endDate", F.to_date("endDate"))

# 2. 生成月度起始日期序列
df = df.withColumn(
    "month_starts",
    F.expr("sequence(date_trunc('month', startDate), date_trunc('month', endDate), interval 1 month)")
)

# 3. 展开序列为多行
df_exploded = df.withColumn("month_start", F.explode("month_starts")).drop("month_starts")

# 4. 计算每个月度子区间的起止日期
df_result = df_exploded.withColumn(
    "startDate", F.greatest(F.col("startDate"), F.col("month_start"))
).withColumn(
    "endDate", F.least(F.col("endDate"), F.last_day(F.col("month_start")))
).select("ID", "startDate", "endDate")

# 查看结果
df_result.show()

关键细节说明 #

• 自动处理同月份区间：若startDate和endDate在同一个月，直接保留原行，无需拆分
• 跨年份兼容：自动识别跨年的日期区间，拆分到对应年份的月份
• 性能优化：PySpark方案通过分布式计算避免单节点内存压力，适合千万级以上数据量

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