PySpark中按分区迭代DataFrame行并保留值计算的实现方案

需求说明 #

需按coll_id_latest分区遍历数据集的每一行，规则如下：

• 每个分区的第一条记录total_alloc初始化为0；
• 其余行根据以下规则生成final_pba列，且total_alloc需保留并用于下一行的计算：
  • 当acct转为小写后为"primary"时：若exposure非空且不为0，且total_alloc < pba1，则final_pba取exposure与(pba1-total_alloc)的较小值，否则为0；
  • 当acct转为小写后为"secondary"且suff_ind转为小写后不等于'y'，且total_alloc < pba1时：final_pba取max(pba1-total_alloc, 0)与exposure的较小值；
• 每次计算后，更新total_alloc为total_alloc + final_pba。

原SAS实现代码 #

data x;
    set y;
    by coll_id_latest;
    retain total_alloc  0;
    if first.coll_id_latest then
        do;total_alloc=0;end;
    if lowcase(acct)="primary" then
                        do;
                            if  exposure not in (.,0) and total_alloc<pba1 then
                                    final_pba   =   min(exposure,(pba1-total_alloc));
                            else    final_pba   =   0;                          
                        end;
    if  lowcase(acct)="secondary" 
                    and lowcase(suff_ind)^='y' and total_alloc < pba1 then
                        do;
                            final_pba   =min(max(pba1-total_alloc,0),exposure);                         
                        end;
                    total_alloc=sum(total_alloc,final_pba);
run;

PySpark高效实现方案 #

由于需求涉及依赖前一行计算结果的累积迭代，逐行循环会导致严重的性能问题和内存溢出。推荐使用Grouped Map Pandas UDF，在每个分区内用Pandas处理顺序逻辑，既保留SAS的retain语义，又能利用Spark的分布式计算能力。

示例代码

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

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

# 定义输出Schema，在输入基础上增加final_pba和total_alloc列
output_schema = StructType([
    StructField("coll_id_latest", StringType(), True),
    StructField("acct", StringType(), True),
    StructField("exposure", DoubleType(), True),
    StructField("pba1", DoubleType(), True),
    StructField("suff_ind", StringType(), True),
    StructField("final_pba", DoubleType(), True),
    StructField("total_alloc", DoubleType(), True)
])

@pandas_udf(output_schema)
def calculate_alloc(df):
    # 确保分组内的行顺序稳定（需提前在Spark中按coll_id_latest和必要字段排序）
    df = df.sort_values("coll_id_latest")
    total_alloc = 0.0
    final_pba_list = []
    total_alloc_list = []
    
    for idx, row in df.iterrows():
        # 分区第一条记录初始化total_alloc
        if idx == 0:
            total_alloc = 0.0
        
        final_pba = 0.0
        acct_lower = str(row["acct"]).lower() if pd.notna(row["acct"]) else ""
        suff_ind_lower = str(row["suff_ind"]).lower() if pd.notna(row["suff_ind"]) else ""
        
        # 处理primary账户逻辑
        if acct_lower == "primary":
            if pd.notna(row["exposure"]) and row["exposure"] != 0 and total_alloc < row["pba1"]:
                final_pba = min(row["exposure"], row["pba1"] - total_alloc)
            else:
                final_pba = 0.0
        # 处理secondary账户逻辑
        elif acct_lower == "secondary" and suff_ind_lower != "y":
            if total_alloc < row["pba1"]:
                final_pba = min(max(row["pba1"] - total_alloc, 0), row["exposure"])
        
        # 更新total_alloc并记录结果
        total_alloc += final_pba
        final_pba_list.append(final_pba)
        total_alloc_list.append(total_alloc)
    
    df["final_pba"] = final_pba_list
    df["total_alloc"] = total_alloc_list
    return df

# 加载输入数据（若未排序，需先执行orderBy保证分组内顺序）
input_df = spark.read.table("your_input_table")
# 未排序时补充：input_df = input_df.orderBy("coll_id_latest")

# 按分区分组计算
result_df = input_df.groupBy("coll_id_latest").apply(calculate_alloc)

# 输出结果
result_df.show()
# result_df.write.saveAsTable("your_output_table")

关键说明

1. 顺序保证：必须确保每个coll_id_latest分组内的行顺序与SAS处理逻辑一致，否则计算结果会出错。若原始数据无固定顺序，需先执行orderBy操作。
2. 性能优势：Grouped Map UDF将每个分组加载到Pandas DataFrame中批量处理，比逐行循环效率提升显著，同时依托Spark分布式架构避免内存溢出。
3. 空值兼容：代码中对acct、suff_ind、exposure的空值做了处理，与SAS逻辑保持一致。

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