PySpark DataFrame按ID分组：保留最早Active行或最早行

需求说明 #

给定如下PySpark DataFrame：

需要按ID分组后执行以下规则生成新DataFrame：

• 若组内存在Active状态的行，保留其中datecreation最早的Active行
• 若组内无Active状态，保留datecreation最早的行

实现方案 #

推荐使用窗口函数实现，逻辑简洁且执行高效，步骤如下：

1. 将datecreation转换为日期类型，确保排序逻辑准确
2. 定义窗口规则：按ID分组，先将Active状态的行排在最前面，再按datecreation升序排序
3. 为每组的行添加行号，取行号为1的行即为目标行

完整代码示例

# 导入必要依赖
from pyspark.sql import Window
from pyspark.sql.functions import col, row_number, when

# 构建测试DataFrame
data = [
    ("A", "nameA1", "2024-12-01", "Active"),
    ("B", "nameB1", "2024-12-10", "Active"),
    ("B", "nameB2", "2024-11-01", "Deleted"),
    ("C", "nameC1", "2024-11-01", "Active"),
    ("C", "nameC2", "2024-10-01", "Active"),
    ("C", "nameC3", "2024-09-02", "Deleted"),
    ("C", "nameC4", "2024-12-03", "Deleted"),
    ("D", "nameD1", "2024-11-01", "Deleted"),
    ("D", "nameD2", "2024-11-02", "Deleted"),
    ("D", "nameD3", "2024-11-03", "Deleted")
]

df = spark.createDataFrame(data, ["ID", "Name", "datecreation", "Status"])

# 转换日期类型
df = df.withColumn("datecreation", col("datecreation").cast("date"))

# 定义窗口：按ID分组，Active行优先，再按日期升序
window_spec = Window.partitionBy("ID").orderBy(
    when(col("Status") == "Active", 0).otherwise(1),
    col("datecreation").asc()
)

# 生成结果
result_df = df.withColumn("row_num", row_number().over(window_spec)) \
              .filter(col("row_num") == 1) \
              .drop("row_num")

# 查看结果
result_df.show()

逻辑说明

• 使用when函数给Active行标记为0，其他行标记为1，排序时0在前，确保Active行优先被选中
• 同一优先级内（同是Active或同是Deleted），按datecreation升序排列，取最早的一行
• 最后过滤行号为1的行，即可得到每组符合规则的目标行

备选方案（聚合方式） #

如果偏好更直观的聚合逻辑，也可以通过分组计算目标日期后筛选：

from pyspark.sql.functions import min

# 计算每个ID是否存在Active行
group_info = df.groupBy("ID").agg(
    (col("Status") == "Active").cast("boolean").alias("has_active")
)

# 关联原数据并筛选目标行
target_df = df.join(group_info, on="ID") \
              .withColumn("target_date", 
                          when(col("has_active"), 
                               min(col("datecreation")).over(Window.partitionBy("ID").where(col("Status") == "Active"))
                              ).otherwise(
                               min(col("datecreation")).over(Window.partitionBy("ID"))
                              )
                         ) \
              .withColumn("is_target", 
                          when(col("has_active"), 
                               (col("Status") == "Active") & (col("datecreation") == col("target_date"))
                              ).otherwise(
                               col("datecreation") == col("target_date")
                              )
                         ) \
              .filter(col("is_target")) \
              .drop("has_active", "target_date", "is_target")

target_df.show()

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