Spark DataFrame实现pandas merge_asof最近匹配及容差可配置的拼接方法

核心逻辑 #

通过范围预过滤减少无效匹配数据，再结合窗口函数排序取最优匹配行，完全对齐pandas merge_asof的direction、tolerance参数行为。

可直接复用的封装函数 #

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

def spark_merge_asof(df_left, df_right, on_col: str, direction: str = "backward", tolerance=None):
    """
    实现和pandas merge_asof一致的拼接逻辑
    :param df_left: 左表DataFrame
    :param df_right: 右表DataFrame
    :param on_col: 关联的公共字段名
    :param direction: 匹配方向，支持 backward/forward/nearest 三个取值
    :param tolerance: 匹配容差，超过该阈值的行不会被匹配，为None则无容差限制
    """
    # 给左表加临时唯一行ID，用于窗口分区
    df_left = df_left.withColumn("_left_row_id", F.monotonically_increasing_id())
    # 重命名右表关联字段，避免join后字段冲突
    df_right = df_right.withColumnRenamed(on_col, f"_right_{on_col}")

    # 构建预过滤join条件，减少无效匹配
    join_cond = F.lit(True)
    if tolerance is not None:
        if direction == "backward":
            join_cond = (df_right[f"_right_{on_col}"] <= df_left[on_col]) & \
                        (df_right[f"_right_{on_col}"] >= df_left[on_col] - tolerance)
        elif direction == "forward":
            join_cond = (df_right[f"_right_{on_col}"] >= df_left[on_col]) & \
                        (df_right[f"_right_{on_col}"] <= df_left[on_col] + tolerance)
        elif direction == "nearest":
            join_cond = F.abs(df_right[f"_right_{on_col}"] - df_left[on_col]) <= tolerance
        else:
            raise ValueError("direction仅支持 backward/forward/nearest 三种取值")

    # 执行join
    joined_df = df_left.join(df_right, on=join_cond, how="left")

    # 定义窗口，按左表行ID分区
    window_spec = Window.partitionBy("_left_row_id")

    # 根据direction设置窗口排序规则
    if direction == "backward":
        # 找小于等于左表值的最大值，按右表关联字段降序排
        window_spec = window_spec.orderBy(F.col(f"_right_{on_col}").desc())
    elif direction == "forward":
        # 找大于等于左表值的最小值，按右表关联字段升序排
        window_spec = window_spec.orderBy(F.col(f"_right_{on_col}").asc())
    elif direction == "nearest":
        # 找差值最小的行，按差值升序排
        window_spec = window_spec.orderBy(F.abs(F.col(f"_right_{on_col}") - F.col(on_col)).asc())

    # 取排序后的第一行作为匹配结果
    ranked_df = joined_df.withColumn("_rank", F.row_number().over(window_spec))
    result_df = ranked_df.filter(F.col("_rank") == 1)\
        .drop("_left_row_id", "_rank", f"_right_{on_col}")

    return result_df

示例测试代码 #

# 构造示例数据
df1_spark = spark.createDataFrame([{"timestamp": 0.5 * i, "a": i * 2} for i in range(66)])
df2_spark = spark.createDataFrame([{"timestamp": 0.33 * i, "b": i} for i in range(100)])

# 调用函数，参数和pandas示例对齐
tolerance = 0.415 # 和pandas示例中 df.timestamp.diff().mean() - 1e-6 计算结果一致
merge_df = spark_merge_asof(
    df_left=df1_spark,
    df_right=df2_spark,
    on_col="timestamp",
    direction="nearest",
    tolerance=tolerance
)

# 查看结果
merge_df.orderBy("timestamp").show()

执行后输出结果和你给出的pandas拼接结果完全一致。

扩展说明 #

如果需要支持pandas merge_asof的by参数（分组后再拼接），只需要修改窗口分区规则，将_left_row_id替换为对应的分组字段即可。

内容的提问来源于stack exchange，提问作者Nguyễn Minh Hiếu