PySpark基于带条件窗口函数实现旅客人数按规则取整

实现思路 #

• 核心逻辑是基于分区窗口计算取整后数值的累计和，再和配额对比赋值结果
• 不需要自定义UDF，直接用Spark内置的窗口函数和数学函数即可完成，性能远高于UDF实现

完整代码示例 #

首先导入依赖：

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

构造示例数据（实际使用时替换为自身数据源即可）：

# 构造示例数据
data = [
    (1, 9, 2.7, 1), (1, 9, 1.3, 2), (1, 9, 1.3, 3),
    (1, 9, 1.3, 4), (1, 9, 1.2, 5), (1, 9, 1.1, 6),
    (2, 9, 2.7, 1), (2, 9, 1.3, 2), (2, 9, 1.3, 3),
    (2, 9, 1.3, 4), (2, 9, 1.2, 5), (2, 9, 1.1, 6)
]
df = spark.createDataFrame(data, schema=["date", "total_trav", "num_trav", "order"])

定义窗口并计算结果列：

# 定义窗口：按date分组，按order升序排序，计算从分组第一行到当前行的累计值
w = Window.partitionBy("date").orderBy("order").rowsBetween(Window.unboundedPreceding, 0)

# 计算结果列
df_result = df.withColumn("ceil_val", F.ceil(F.col("num_trav"))) \
              .withColumn("cum_sum", F.sum("ceil_val").over(w)) \
              .withColumn("trav_res", F.when(F.col("cum_sum") <= F.col("total_trav"), F.col("ceil_val")).otherwise(0)) \
              .drop("ceil_val", "cum_sum") # 删除中间临时列

验证输出：

df_result.orderBy("date", "order").show()

输出结果和你给出的预期示例完全一致。

逻辑说明 #

1. 首先用ceil函数对num_trav做向上取整得到临时列ceil_val
2. 按date分区、order升序的窗口，计算从分组首行到当前行的ceil_val累计和cum_sum
3. 对比累计和与当前分组的配额total_trav，累计和未超过配额则结果取当前行的ceil_val，超过则取0

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