PySpark中确保数组相邻元素间隔≥3天及关联数组同步处理

这个需求的核心是每行独立处理时间戳数组，既要保证保留的startTimeArray相邻元素时间差至少3天，又要同步调整endTimeArray，同时尽可能减少删除操作。因为PySpark内置函数很难处理这种带依赖的数组过滤逻辑，所以用自定义UDF（User Defined Function）是最直接的方式。

核心逻辑拆解 #

先明确规则对应的操作步骤：

1. 优先保留startTimeArray的第一个元素，从第二个元素开始，只保留与上一个保留元素时间差≥3天的项；
2. 每删除一个startTimeArray中索引为i的元素，就要同步删除endTimeArray中索引为i-1的元素；
3. 如果数组长度≤1，直接原样保留。

实现代码 #

1. 导入依赖模块

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, TimestampType, StructType, StructField
import datetime

2. 定义处理数组的逻辑函数

这个函数会接收每行的startTimeArray和endTimeArray，返回处理后的新数组：

def process_time_arrays(start_arr, end_arr):
    # 数组长度≤1时直接返回
    if len(start_arr) <= 1:
        return (start_arr, end_arr)
    
    # 第一步：筛选符合条件的start索引（保留第一个，之后只留与前一个保留项差≥3天的）
    kept_start_indices = [0]
    last_kept_start = start_arr[0]
    
    for i in range(1, len(start_arr)):
        current_start = start_arr[i]
        time_diff = current_start - last_kept_start
        if time_diff.days >= 3:
            kept_start_indices.append(i)
            last_kept_start = current_start
    
    # 第二步：确定要删除的start索引
    deleted_start_indices = set(range(len(start_arr))) - set(kept_start_indices)
    
    # 第三步：确定要删除的end索引（对应start索引i的i-1，且i-1≥0）
    deleted_end_indices = set()
    for i in deleted_start_indices:
        if i - 1 >= 0:
            deleted_end_indices.add(i-1)
    
    # 第四步：构建新的start和end数组
    new_start = [start_arr[i] for i in kept_start_indices]
    new_end = [end_arr[j] for j in range(len(end_arr)) if j not in deleted_end_indices]
    
    return (new_start, new_end)

3. 注册UDF并应用到DataFrame

因为UDF要返回两个数组，所以用StructType定义返回类型：

# 注册UDF
process_arrays_udf = udf(
    process_time_arrays,
    returnType=StructType([
        StructField("newStartTimeArray", ArrayType(TimestampType())),
        StructField("newEndTimeArray", ArrayType(TimestampType()))
    ])
)

# 应用到原始DataFrame（假设原始DF名为df）
result_df = df.withColumn("processed", process_arrays_udf("startTimeArray", "endTimeArray")) \
              .select(
                  "category",
                  "processed.newStartTimeArray",
                  "processed.newEndTimeArray"
              ) \
              .withColumnRenamed("newStartTimeArray", "startTimeArray") \
              .withColumnRenamed("newEndTimeArray", "endTimeArray")

验证示例 #

拿你提供的测试数据来看：

• 第一行原始startTimeArray是[2019-01-10, 2019-01-12, 2019-01-16, 2019-01-20]，处理后保留[2019-01-10, 2019-01-16, 2019-01-20]；对应的endTimeArray删除了索引0的元素，得到[2019-01-15, 2019-01-18, 2019-01-22]，完全匹配示例输出。
• 第二行和第三行的处理结果也和示例一致，说明逻辑是正确的。

边界情况说明 #

• 当startTimeArray只有一个元素时，直接保留原数组；
• 当连续多个start元素都不符合时间差要求时，只会保留第一个和后续符合条件的项，对应的end数组会同步删除对应的前置元素。

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