Spark DataFrame按caseid分组优先保留non-STP记录实现咨询

核心逻辑：给同caseid下的记录设置优先级，non-STP优先级高于STP，仅保留每个caseid优先级最高的一条记录即可。

方法1：窗口函数实现（推荐） #

性能更优，仅需一次shuffle操作，代码简洁：

PySpark 示例代码 #

from pyspark.sql import Window
from pyspark.sql.functions import row_number, when

# 替换为你的原始DataFrame变量名
raw_df = <你的原始DataFrame>

# 按caseid分组，同组内non-STP排在最前
window_rule = Window.partitionBy("caseid").orderBy(
    when(raw_df.indicator == "non-STP", 1).otherwise(2)
)

result_df = raw_df.withColumn("rank", row_number().over(window_rule)) \
    .filter("rank = 1") \
    .drop("rank")

Scala 示例代码 #

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions.{row_number, when}

// 替换为你的原始DataFrame变量名
val rawDf = <你的原始DataFrame>

val windowRule = Window.partitionBy("caseid").orderBy(
  when($"indicator" === "non-STP", 1).otherwise(2)
)

val resultDf = rawDf.withColumn("rank", row_number.over(windowRule))
  .filter($"rank" === 1)
  .drop("rank")

方法2：分组聚合关联实现 #

逻辑更直观，适合新手理解：

PySpark 示例代码 #

from pyspark.sql.functions import max, when

raw_df = <你的原始DataFrame>

# 先统计每个caseid最终要保留的indicator值
caseid_target_df = raw_df.groupBy("caseid").agg(
    max(when(raw_df.indicator == "non-STP", "non-STP").otherwise("STP")).alias("target_ind")
)

# 关联原表过滤得到符合要求的记录
result_df = raw_df.join(caseid_target_df, on="caseid", how="inner") \
    .filter(raw_df.indicator == caseid_target_df.target_ind) \
    .select("caseid", "indicator")

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