Scala Spark中基于窗口分区的动态条件累积求和实现购买量封顶

针对你的需求——按用户、日期排序，每3天窗口内最多允许3次购买，计算每日的cappedPurchases，以下是两种高性能的实现方案，替代性能较差的collect_list UDF：

方案一：Pandas UDF（推荐，大数据量友好） #

Pandas UDF基于向量化处理，比普通Python UDF性能提升显著，且分组处理逻辑清晰，避免了全量数据拉取的问题。

实现代码（PySpark） #

from pyspark.sql import SparkSession
from pyspark.sql.functions import pandas_udf, col
from pyspark.sql.types import StructType, StructField, StringType, IntegerType
import pandas as pd

# 定义处理每个用户分组的逻辑
def calculate_capped_purchases(user_df: pd.DataFrame) -> pd.DataFrame:
    # 确保数据按日期升序排列
    user_df = user_df.sort_values("day").reset_index(drop=True)
    
    capped_list = []
    used_quota = 0
    # 用队列维护最近3天的capped值，实时计算窗口内已用额度
    window_queue = []
    
    for _, row in user_df.iterrows():
        current_day = row["day"]
        current_num = row["numPurchases"]
        
        # 移除窗口中超出[当前日-2, 当前日]的历史记录
        while window_queue and window_queue[0]["day"] < current_day - 2:
            removed = window_queue.pop(0)
            used_quota -= removed["capped"]
        
        # 计算当日可用额度
        available = 3 - used_quota
        if available <= 0:
            current_capped = 0
        else:
            current_capped = min(current_num, available)
        
        capped_list.append(current_capped)
        # 更新已用额度和窗口队列
        used_quota += current_capped
        window_queue.append({"day": current_day, "capped": current_capped})
    
    user_df["cappedPurchases"] = capped_list
    return user_df

# 定义输出Schema
output_schema = StructType([
    StructField("user", StringType()),
    StructField("day", IntegerType()),
    StructField("numPurchases", IntegerType()),
    StructField("cappedPurchases", IntegerType())
])

# 注册Pandas UDF
capped_udf = pandas_udf(calculate_capped_purchases, output_schema)

# 应用到数据集（假设你的原始DataFrame名为raw_df）
result_df = raw_df.groupBy("user").apply(capped_udf)
result_df.show()

性能优势 #

• 向量化处理：Pandas UDF利用Apache Arrow传递数据，减少序列化开销，比普通Python UDF快数倍
• 分组流式处理：每个用户的数据独立处理，不会将全量数据加载到内存
• 时间复杂度O(n)：每个用户分组内仅需遍历一次数据，队列操作保证滑动窗口计算高效

方案二：递归CTE（SQL环境适用） #

如果你的场景依赖SQL处理，可以用递归CTE实现逐天的额度累积计算，适合日期连续的数据集。

实现SQL #

WITH recursive user_purchase_sequence AS (
    -- 初始行：每个用户的最早日期
    SELECT
        user,
        day,
        numPurchases,
        LEAST(numPurchases, 3) AS cappedPurchases,
        LEAST(numPurchases, 3) AS used_in_window
    FROM your_table
    WHERE day = (SELECT MIN(day) FROM your_table WHERE user = your_table.user)
    
    UNION ALL
    
    -- 递归计算后续日期
    SELECT
        t.user,
        t.day,
        t.numPurchases,
        -- 计算当日可购买额度
        CASE
            WHEN (r.used_in_window - COALESCE(p3.cappedPurchases, 0)) >= 3 THEN 0
            ELSE LEAST(t.numPurchases, 3 - (r.used_in_window - COALESCE(p3.cappedPurchases, 0)))
        END AS cappedPurchases,
        -- 更新最近3天的已用额度：减去3天前的capped值，加上当日的capped值
        (r.used_in_window - COALESCE(p3.cappedPurchases, 0)) +
        CASE
            WHEN (r.used_in_window - COALESCE(p3.cappedPurchases, 0)) >= 3 THEN 0
            ELSE LEAST(t.numPurchases, 3 - (r.used_in_window - COALESCE(p3.cappedPurchases, 0)))
        END AS used_in_window
    FROM your_table t
    JOIN user_purchase_sequence r 
        ON t.user = r.user AND t.day = r.day + 1
    -- 关联3天前的capped记录（用于更新窗口额度）
    LEFT JOIN user_purchase_sequence p3
        ON t.user = p3.user AND p3.day = t.day - 3
)
SELECT user, day, numPurchases, cappedPurchases
FROM user_purchase_sequence
ORDER BY user, day;

注意事项 #

• 要求用户的日期是连续的（无缺失天），否则递归会中断
• 对于日期跨度极大的用户分组，性能可能不如Pandas UDF

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