Scala Spark如何将窗口分区数据集传入UDF实现分区打标逻辑

你可以选择以下两种方案实现需求，优先推荐第一种窗口函数方案，性能更高且不需要自定义函数处理全分区数据：

方案1：使用Spark窗口函数实现 #

你不需要自定义UDF/UDAF，直接通过同分区全局聚合的窗口函数计算所需的两个对比指标即可：

1. 按Id2、Id3分区，计算分区内TaskId = "hu77d9-mnb"对应的行的index值
2. 同分区计算TaskName包含deploy的行的最小index值（如果有多个含deploy的任务，取最早出现的即可，可根据业务调整取值逻辑）
3. 对比两个指标值给所有行打result标记

SQL写法示例： #

WITH partition_metrics AS (
    SELECT
        *,
        -- 取当前分区内目标任务的index
        MAX(CASE WHEN TaskId = 'hu77d9-mnb' THEN index ELSE NULL END) OVER (PARTITION BY Id2, Id3) AS target_task_index,
        -- 取当前分区内首个含deploy任务的index
        MIN(CASE WHEN TaskName LIKE '%deploy%' THEN index ELSE NULL END) OVER (PARTITION BY Id2, Id3) AS first_deploy_index
    FROM 你的原始表名
)
SELECT
    *,
    -- 符合条件打1，否则打0，可自行调整边界场景（比如缺目标任务/缺deploy任务时的取值）
    CASE 
        WHEN target_task_index IS NOT NULL 
            AND first_deploy_index IS NOT NULL 
            AND target_task_index < first_deploy_index 
        THEN 1 
        ELSE 0 
    END AS result
FROM partition_metrics

方案2：使用mapPartitions处理全分区数据 #

如果你的后续逻辑还有更复杂的分区级处理要求，可以用mapPartitions算子直接拿到整个分区的所有行处理：

注意：使用该方案前需要确保按Id2、Id3重新分区，保证同组数据落在同一个分区内，避免逻辑错误。如果单分区数据量过大，全量加载到列表可能有OOM风险。

PySpark写法示例： #

from pyspark.sql import Row

def process_single_partition(partition_rows):
    row_list = list(partition_rows)
    target_index = None
    deploy_index_list = []
    
    # 遍历分区内所有行，提取所需指标
    for row in row_list:
        if row.TaskId == "hu77d9-mnb":
            target_index = row.index
        if "deploy" in row.TaskName:
            deploy_index_list.append(row.index)
    
    # 计算当前分区的result值
    current_result = 0
    if target_index is not None and len(deploy_index_list) > 0:
        min_deploy_index = min(deploy_index_list)
        if target_index < min_deploy_index:
            current_result = 1
    
    # 给分区内所有行添加result字段后返回
    for row in row_list:
        row_dict = row.asDict()
        row_dict["result"] = current_result
        yield Row(**row_dict)

# 调用算子生成结果表
result_df = df.repartition("Id2", "Id3").rdd.mapPartitions(process_single_partition).toDF()

边界情况说明 #

如果分区内不存在TaskId = "hu77d9-mnb"的行，或者不存在包含deploy的TaskName，可以根据业务需要调整result的默认取值，上述示例中默认赋值为0。

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