Spark：计算乘客从印度出境至返回前的最长旅行连续航班数

核心思路 #

通过Window函数实现分布式的行程区间划分与统计，避免collect_list()+UDF带来的内存瓶颈与性能损耗，步骤如下：

实现步骤 #

1. 转换日期格式（确保行程顺序正确）

原数据中date为字符串类型，需转换为日期类型保证排序逻辑准确：

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

# 转换日期列格式
df = df.withColumn("date", F.to_date(F.col("date"), "dd/MM/yyyy"))

2. 划分出境-入境行程组

按乘客分组、日期排序，为每个"从印度出境"事件标记独立的行程组ID，后续航班自动归属到当前组，直到下一次出境事件：

# 定义乘客级别的排序窗口
window_passenger = Window.partitionBy("passengerId").orderBy("date")

# 生成行程组ID：每次遇到出境（indOriginDestFlg=1）时组ID递增
df = df.withColumn(
    "trip_group",
    F.sum(F.when(F.col("indOriginDestFlg") == 1, 1).otherwise(0)).over(window_passenger)
)

3. 过滤无返回印度的无效组

仅保留"出境后有返回印度"的行程组，排除那些出境后未返回的无效区间：

# 标记每个行程组是否存在返回印度的记录
group_has_return = df.groupBy("passengerId", "trip_group").agg(
    F.max(F.when(F.col("indOriginDestFlg") == 2, 1).otherwise(0)).alias("has_return")
)

# 关联回原数据
df = df.join(group_has_return, on=["passengerId", "trip_group"], how="left")

4. 计算每组的连续航班数并取最大值

有效组内的总航班数减去出境、入境两个航班，得到中间连续旅行的航班数，最后按乘客取最大值：

# 计算每个有效组的连续航班数
df = df.withColumn(
    "streak",
    F.when(
        F.col("has_return") == 1,
        F.count("*").over(Window.partitionBy("passengerId", "trip_group")) - 2
    ).otherwise(0)
)

# 提取每个乘客的最长连续航班数
result = df.groupBy("passengerId").agg(
    F.max(F.col("streak")).alias("maxTravelStreak")
)

# 查看结果
result.show()

方案优势 #

• 分布式处理：全程基于Spark分布式计算框架，避免collect_list()将数据拉取到Driver端的内存压力，支持大规模数据场景。
• 无UDF依赖：无需编写自定义UDF，减少代码复杂度与序列化开销。
• 逻辑清晰：通过行程组划分直观对应"出境-返回"的完整区间，便于维护与扩展。

结果验证 #

运行后将得到与预期一致的输出：

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