PySpark实现：在verrekenbaar间隙添加niet-verrekenbaar时间行

核心逻辑是verrekenbaar类型时间范围优先保留，将niet-verrekenbaar的时间范围按verrekenbaar的分段切割，自动填充间隙部分为新的niet-verrekenbaar行，最终所有时间分段连续且无重叠。

输入示例 #

data1 = [
    (123, 202009, 1, 'niet-verrekenbaar', '2020-09-13T00:00:00.000Z', '2020-09-25T00:00:00.000Z'),
    (123, 202009, 1, 'verrekenbaar', '2020-09-16T00:00:00.000Z', '2020-09-17T00:00:00.000Z'),
    (123, 202009, 1, 'verrekenbaar', '2020-09-19T00:00:00.000Z', '2020-09-20T00:00:00.000Z'),
    (123, 202009, 1, 'verrekenbaar', '2020-09-22T00:00:00.000Z', '2020-09-27T00:00:00.000Z'),
    (123, 202009, 1, 'niet-verrekenbaar', '2020-09-26T00:00:00.000Z', '2020-09-31T00:00:00.000Z'),
    (123, 202009, 1, 'verrekenbaar', '2020-09-29T00:00:00.000Z', '2020-09-31T00:00:00.000Z')
]

预期输出 #

(123, 202009, 1, 'niet-verrekenbaar', '2020-09-13T00:00:00.000Z', '2020-09-16T00:00:00.000Z'),
(123, 202009, 1, 'verrekenbaar', '2020-09-16T00:00:00.000Z', '2020-09-17T00:00:00.000Z'),
(123, 202009, 1, 'niet-verrekenbaar', '2020-09-17T00:00:00.000Z', '2020-09-19T00:00:00.000Z'),
(123, 202009, 1, 'verrekenbaar', '2020-09-19T00:00:00.000Z', '2020-09-20T00:00:00.000Z'),
(123, 202009, 1, 'niet-verrekenbaar', '2020-09-20T00:00:00.000Z', '2020-09-22T00:00:00.000Z'),
(123, 202009, 1, 'verrekenbaar', '2020-09-22T00:00:00.000Z', '2020-09-27T00:00:00.000Z'),
(123, 202009, 1, 'niet-verrekenbaar', '2020-09-27T00:00:00.000Z', '2020-09-29T00:00:00.000Z'),
(123, 202009, 1, 'verrekenbaar', '2020-09-29T00:00:00.000Z', '2020-09-31T00:00:00.000Z')

PySpark实现代码 #

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.types import StructType, StructField, IntegerType, StringType, TimestampType

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

# 定义数据Schema
schema = StructType([
    StructField("user_id", IntegerType(), True),
    StructField("month", IntegerType(), True),
    StructField("type_id", IntegerType(), True),
    StructField("status", StringType(), True),
    StructField("start_time", TimestampType(), True),
    StructField("end_time", TimestampType(), True)
])

# 创建原始DataFrame
df = spark.createDataFrame(data1, schema=schema)

# 分组处理生成最终结果
result_df = df.groupBy("user_id", "month", "type_id").agg(
    # 收集所有起止时间点，合并为列表
    F.flatten(F.array(
        F.collect_list(F.array("start_time")),
        F.collect_list(F.array("end_time"))
    )).alias("time_points")
).withColumn(
    # 对时间点排序，生成有序刻度
    "sorted_times", F.sort_array("time_points")
).withColumn(
    # 用相邻时间点生成连续的时间分段
    "time_ranges", F.expr("transform(slice(sorted_times, 1, size(sorted_times)-1), (x, i) -> struct(x as start_time, sorted_times[i+1] as end_time))")
).withColumn(
    # 展开所有分段
    "time_range", F.explode("time_ranges")
).select(
    "user_id", "month", "type_id",
    F.col("time_range.start_time").alias("start_time"),
    F.col("time_range.end_time").alias("end_time")
).withColumn(
    # 标记分段状态：优先匹配verrekenbaar，其余为niet-verrekenbaar
    "status", F.when(
        F.exists(
            df.select("user_id", "month", "type_id", "status", "start_time", "end_time"),
            lambda r: (
                r.user_id == F.col("user_id")
                & r.month == F.col("month")
                & r.type_id == F.col("type_id")
                & r.status == "verrekenbaar"
                & r.start_time <= F.col("start_time")
                & r.end_time >= F.col("end_time")
            )
        ),
        "verrekenbaar"
    ).otherwise("niet-verrekenbaar")
).filter(
    # 过滤掉起止时间相同的无效分段
    F.col("start_time") != F.col("end_time")
).orderBy(
    # 按用户、月份、类型、时间排序
    "user_id", "month", "type_id", "start_time"
)

# 查看结果
result_df.show(truncate=False)

# 若需要转换为原始元组格式，可执行以下代码
# result_tuples = result_df.rdd.map(tuple).collect()

关键逻辑说明 #

1. 收集排序时间节点：把同一用户、月份、类型下的所有时间起止点收集并排序，得到连续的时间刻度
2. 生成连续分段：用相邻时间刻度组成小的时间区间，确保所有区间连续无重叠
3. 标记分段状态：检查每个分段是否被verrekenbaar的时间范围完全覆盖，是则标记对应状态，剩余间隙自动设为niet-verrekenbaar
4. 过滤无效分段：剔除起止时间相同的空区间，避免无效数据

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