PySpark如何不使用UDF高效合并重叠区间并记录合并ID列表

完全可以不使用UDF实现该需求，仅依靠PySpark内置窗口函数即可完成，适配大数据量场景，性能远高于自定义UDF方案。

实现逻辑 #

• 首先对所有区间按start字段升序排序，不需要依赖原数据的id或start字段的原有顺序
• 通过窗口函数计算截至当前行之前所有区间的最大end值
• 对区间做分组标记：如果当前行的start大于前面所有区间的最大end，说明是新的不重叠区间，分配新分组ID，否则归属到上一个分组
• 按分组ID聚合，取每组最小start、最大end，同时收集所有原区间的id得到ids字段

完整代码示例 #

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

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

# 构造示例数据，实际场景替换为你自己的DataFrame即可
data = [(0,10,20),(1,11,13),(2,14,18),(3,22,30),(4,25,27),(5,28,31)]
df = spark.createDataFrame(data, schema=["id", "start", "end"])

# 定义按start排序的窗口
w_order = Window.orderBy("start")
# 计算当前行之前所有区间的最大end值
df = df.withColumn("prev_max_end", F.max("end").over(w_order.rowsBetween(Window.unboundedPreceding, -1)))
# 标记是否为新的不重叠区间
df = df.withColumn("new_group", F.when(F.col("prev_max_end").isNull() | (F.col("start") > F.col("prev_max_end")), 1).otherwise(0))
# 累加标记得到分组ID
w_group = Window.orderBy("start").rowsBetween(Window.unboundedPreceding, Window.currentRow)
df = df.withColumn("group_id", F.sum("new_group").over(w_group))

# 按分组聚合得到最终结果
result = df.groupBy("group_id") \
           .agg(
               F.min("start").alias("start"),
               F.max("end").alias("end"),
               F.collect_list("id").alias("ids")
           ) \
           .drop("group_id")

result.show()

输出结果 #

+-----+---+---------+
|start|end|      ids|
+-----+---+---------+
|   10| 20|[0, 1, 2]|
|   22| 31|[3, 4, 5]|
+-----+---+---------+

所有计算逻辑均使用Spark原生内置函数，避免了UDF的序列化反序列化开销，适合TB级以上的大数据量场景。

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