PySpark中如何实现类SAS Retain的逐行依赖计算？

需求说明 #

需要在PySpark中实现类似SAS Retain语句的逐行计算逻辑：按顺序遍历表时，当前行使用前一行计算后的值（而非原始值）。使用Window函数的lag方法只能获取前一行的原始值，无法传递更新后的状态，因此需要其他解决方案。

具体需求：

• 初始行的prev_start为NaN；若前一行prev_start缺失，当前行使用自身start值（如示例第二行prev_start设为1.0）
• 后续行的prev_start继承前一行更新后的prev_start值（如示例第三、四、五行的prev_start均为1.0）
• 同时实现prev_end = max(prev_end, stop)的逻辑，当prev_end小于当前行start时，识别为非连续时间帧

当前测试代码及结果 #

测试代码 #

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

# Create a sample DataFrame
data = [
    (1, 1, 3),
    (2, 2, 4),
    (3, 3, 5),
    (4, 4, 6),
    (5, 8, 10)
]

columns = ["id", "start", "stop"]

spark = SparkSession.builder.master("local").appName("SelfJoin").getOrCreate()

df = spark.createDataFrame(data, columns).sort('start')

# 1. Sort the DataFrame by 'start'
df = df.sort("start").withColumn('prev_start', F.col('start'))

# 2. Initialize a window that looks back one record
window = Window.orderBy(['start']).rowsBetween(-1, -1)

df = (
    df
    .withColumn("prev_start", F.lag("prev_start", 1).over(window))
)

df.toPandas().style.hide_index()

当前结果 #

期望结果 #

解决方案 #

由于Window函数无法跟踪递推更新的状态，我们可以通过两种方式实现需求：迭代式Pandas UDF（适合小数据量/测试场景）或flatMapGroupsWithState（适合大数据量分布式场景）。

方法一：迭代式Pandas UDF #

利用Pandas的逐行迭代特性维护状态，实现递推计算：

from pyspark.sql import SparkSession
from pyspark.sql.functions import pandas_udf, col
from pyspark.sql.types import StructType, StructField, IntegerType, DoubleType

# 初始化SparkSession
spark = SparkSession.builder.master("local").appName("RetainLogic").getOrCreate()

# 构建示例数据
data = [
    (1, 1, 3),
    (2, 2, 4),
    (3, 3, 5),
    (4, 4, 6),
    (5, 8, 10)
]
columns = ["id", "start", "stop"]
df = spark.createDataFrame(data, columns).sort("start")

# 定义输出Schema，包含额外的prev_end和连续状态标识
output_schema = StructType([
    StructField("id", IntegerType(), True),
    StructField("start", IntegerType(), True),
    StructField("stop", IntegerType(), True),
    StructField("prev_start", DoubleType(), True),
    StructField("prev_end", DoubleType(), True),
    StructField("is_continuous", IntegerType(), True)
])

# 定义迭代式pandas UDF，维护递推状态
@pandas_udf(output_schema)
def retain_logic(iterator):
    prev_start = None
    prev_end = None
    for df_batch in iterator:
        result_rows = []
        for _, row in df_batch.iterrows():
            current_start = row['start']
            current_stop = row['stop']
            
            # 处理prev_start和状态更新
            if prev_start is None:
                current_prev_start = None
                prev_start = current_start
                prev_end = current_stop
                is_continuous = 1
            else:
                current_prev_start = prev_start
                # 判断是否连续并更新prev_end
                if prev_end >= current_start:
                    prev_end = max(prev_end, current_stop)
                    is_continuous = 1
                else:
                    prev_start = current_start
                    prev_end = current_stop
                    is_continuous = 0
            
            result_rows.append((row['id'], current_start, current_stop, current_prev_start, prev_end, is_continuous))
        
        # 将结果转换为Pandas DataFrame返回
        import pandas as pd
        yield pd.DataFrame(result_rows, columns=output_schema.fieldNames())

# 应用UDF并查看结果
result_df = df.groupBy().apply(retain_logic)
result_df.show()

方法二：flatMapGroupsWithState（分布式场景） #

通过虚拟分组维护全局状态，适合大数据量的分布式处理：

from pyspark.sql import SparkSession
import pyspark.sql.functions as F
from pyspark.sql.types import StructType, StructField, IntegerType, DoubleType
from pyspark.sql.streaming import GroupState

# 初始化SparkSession
spark = SparkSession.builder.master("local").appName("RetainState").getOrCreate()

# 构建示例数据
data = [
    (1, 1, 3),
    (2, 2, 4),
    (3, 3, 5),
    (4, 4, 6),
    (5, 8, 10)
]
columns = ["id", "start", "stop"]
df = spark.createDataFrame(data, columns).sort("start")

# 添加虚拟分组键，将所有数据归为一组
df = df.withColumn("group_key", F.lit(1))

# 定义状态Schema，存储递推的prev_start和prev_end
state_schema = StructType([
    StructField("prev_start", DoubleType(), True),
    StructField("prev_end", DoubleType(), True)
])

# 定义输出Schema
output_schema = StructType([
    StructField("id", IntegerType(), True),
    StructField("start", IntegerType(), True),
    StructField("stop", IntegerType(), True),
    StructField("prev_start", DoubleType(), True),
    StructField("prev_end", DoubleType(), True),
    StructField("is_continuous", IntegerType(), True)
])

# 定义状态更新函数，逐行处理并维护状态
def update_state(key, rows, state: GroupState):
    # 初始化状态
    if not state.exists:
        prev_start = None
        prev_end = None
    else:
        state_data = state.get()
        prev_start = state_data["prev_start"]
        prev_end = state_data["prev_end"]
    
    # 确保行按start排序
    sorted_rows = sorted(rows, key=lambda x: x["start"])
    
    for row in sorted_rows:
        current_start = row["start"]
        current_stop = row["stop"]
        
        if prev_start is None:
            current_prev_start = None
            prev_start = current_start
            prev_end = current_stop
            is_continuous = 1
        else:
            current_prev_start = prev_start
            if prev_end >= current_start:
                prev_end = max(prev_end, current_stop)
                is_continuous = 1
            else:
                prev_start = current_start
                prev_end = current_stop
                is_continuous = 0
        
        # 输出当前行结果
        yield (row["id"], current_start, current_stop, current_prev_start, prev_end, is_continuous)
    
    # 更新全局状态
    state.update({"prev_start": prev_start, "prev_end": prev_end})

# 应用状态更新函数并去除虚拟分组键
result_df = df.groupBy("group_key").flatMapGroupsWithState(
    outputMode="append",
    stateSchema=state_schema,
    func=update_state
).drop("group_key")

result_df.show()

关键说明 #

Window函数的lag方法仅能访问原始行的静态值，无法跟踪递推过程中动态更新的状态。上述两种方法通过维护全局状态并逐行迭代处理，完美实现了SAS Retain语句的核心逻辑。

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