PySpark如何仅移除batch_date连续月份的同hash重复行

经过多次按batch_date迭代和关联操作后，得到如下DataFrame，我们重点关注r_id0==0和r_id0==1的分组：

+-----+----------+----------+----------+--------------------+--------------------+
|r_id0|batch_date|      from|     until|                desc|                hash|
+-----+----------+----------+----------+--------------------+--------------------+
|    0|2020-01-31|2020-01-31|2020-02-27|  Esto es un campo 1|b10192333e9c61a40...|
|    0|2020-04-30|2020-04-30|9999-12-31|  Esto es un campo 1|b10192333e9c61a40...|
|    1|2020-01-31|2020-01-31|9999-12-31|  Esto es un campo 1|b63b3e8201cd417bb...|
|    1|2020-02-28|2020-02-28|9999-12-31|  Esto es un campo 1|b63b3e8201cd417bb...|
|    1|2020-03-31|2020-03-31|9999-12-31|  Esto es un campo 1|b63b3e8201cd417bb...|
|    2|2020-01-31|2020-01-31|2020-02-27|  Esto es un campo 2|b31c138a4b6a96169...|
|    2|2020-02-28|2020-02-28|2020-03-30|  Esto es un campo 4|fa1cfe1eaf9b14f88...|
|    2|2020-03-31|2020-03-31|2020-04-29|  Esto es un campo 3|e4ffd0cd4a6cf1193...|
|    2|2020-04-30|2020-04-30|9999-12-31| Esto es un campo 23|0fdb24b8fcf8603ee...|
|    3|2020-02-28|2020-02-28|2020-03-30|  Esto es un campo 3|a3a6870ca9b42ad06...|
|    3|2020-03-31|2020-03-31|2020-04-29|  Esto es un campo 2|18b24f88271e99618...|
|    4|2020-03-31|2020-03-31|2020-04-29|  Esto es un campo 5|2fe50db0156cfc909...|
|    6|2020-01-31|2020-01-31|2020-02-27|  Esto es un campo 6|7c6d329b73d9de59f...|
|    6|2020-04-30|2020-04-30|9999-12-31|Esto es un campo 77|d70c9340f83167e95...|
+-----+----------+----------+----------+--------------------+--------------------+

删除hash列相同的行，但仅删除batch_date列对应月份连续的重复行。符合该条件的重复组取首次出现行的全部字段值，以及组内最后一行的until列值；非连续的重复行需保留原状。

示例说明 #

• r_id0==0的两条记录batch_date间隔3个月，需原样保留
• r_id0==1的三条记录batch_date月份连续，需合并为1行，预期结果如下：

+-----+----------+----------+----------+--------------------+--------------------+
|r_id0|batch_date|      from|     until|                desc|                hash|
+-----+----------+----------+----------+--------------------+--------------------+
|    0|2020-01-31|2020-01-31|2020-02-27|  Esto es un campo 42|b10192333e9c61a40...|
|    0|2020-04-30|2020-04-30|9999-12-31|  Esto es un campo 42|b10192333e9c61a40...|
|    1|2020-01-31|2020-01-31|9999-12-31|  Esto es un campo 1|b10192333e9c61a40...|
|    2|2020-01-31|2020-01-31|2020-02-27|  Esto es un campo 2|b31c138a4b6a96169...|
etc etc... 

当前编写的函数可以处理r_id01的场景，但会错误合并r_id00的记录，代码如下：

def removing_duplicates2(final_df):
    '''
    Removing duplicates. It takes all from first time they appear and last "until" column.
    '''
    # we create this list in order to avor_id iteration over the column "hash", which will be our groupby column
    iterating_columns= final_df.columns
    iterating_columns.remove("hash")
    exprs =  [F.first(x).alias(x) if x!="until" else F.last(x).alias(x)  for x in iterating_columns] 
    return  (
         final_df.groupBy("hash").agg(*exprs)
            .dropDuplicates(["hash"]).select(columns)
         )

错误输出结果如下，r_id0==0的两条记录被误合并：

+-----+----------+----------+----------+--------------------+--------------------+
|r_id0|batch_date|      from|     until|                desc|                hash|
+-----+----------+----------+----------+--------------------+--------------------+
|    0|2020-01-31|2020-01-31|9999-12-31|  Esto es un campo 42|b10192333e9c61a40...|
|    1|2020-01-31|2020-01-31|9999-12-31|  Esto es un campo 1|b63b3e8201cd417bb...|
|    2|2020-01-31|2020-01-31|2020-02-27|  Esto es un campo 2|b31c138a4b6a96169...|
etc etc

要求：使用PySpark实现该需求，逻辑要应用到按batch_date分区的超大规模表上，已有遍历batch_date的循环，不希望新增更多循环以免引发性能问题。 #

核心思路是用Spark原生窗口函数实现连续区间标记，完全避免自定义循环，利用Spark分布式计算能力适配超大规模表，实现步骤如下：

1. 对相同hash的记录按batch_date升序排序，计算当前行和上一行batch_date的月份差
2. 对每个hash分组，月份差大于1的位置标记为新分组的起点，累加生成子组ID
3. 按hash+子组ID分组聚合，除until取组内最后一个值外，其余字段取组内第一个值即可
   实现代码如下：

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

def removing_duplicates_correct(final_df):
    # 定义窗口：按hash分区，batch_date升序排序
    w = Window.partitionBy("hash").orderBy("batch_date")
    # 步骤1：计算上一行的batch_date，判断月份是否连续
    df_with_prev = final_df.withColumn("prev_batch", F.lag("batch_date").over(w)) \
        .withColumn("month_diff", F.months_between("batch_date", "prev_batch")) \
        .fillna(0, subset=["month_diff"]) # 组内第一行没有上一行，差值设为0
    # 步骤2：生成连续区间子组ID，差值>1则为新的子组
    df_with_group = df_with_prev.withColumn("new_group_flag", F.when(F.col("month_diff") > 1, 1).otherwise(0)) \
        .withColumn("sub_group_id", F.sum("new_group_flag").over(w.rangeBetween(Window.unboundedPreceding, 0)))
    # 步骤3：按hash+子组ID聚合
    iter_cols = [c for c in final_df.columns if c != "hash"]
    agg_exprs = [F.first(c).alias(c) if c != "until" else F.last(c).alias(c) for c in iter_cols]
    result_df = df_with_group.groupBy("hash", "sub_group_id").agg(*agg_exprs) \
        .select(final_df.columns) # 按原字段顺序输出
    return result_df

该方案所有逻辑均为Spark原生算子，没有额外循环，完全利用集群分布式计算资源，适配超大规模分区表场景。 #

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