PySpark中基于区间数组匹配的DataFrame内连接实现问题

我有如下两个DataFrame，需要实现内连接，连接条件为df2的cep字段落在df的cep_ori数组中任意一个(inicial, final)区间内。

df数据及Schema #

数据： #

state   interval      limits       cep_ori
PA     >=5, <=10      >10    [{111, 222}, {333, 444}]
SP     >=6, <=10      <8     [{333, 444}, {555, 666}]

Schema： #

root
 |-- state: string (nullable = true)
 |-- interval: string (nullable = true)
 |-- limits: string (nullable = true)
 |-- cep_ori: array (nullable = false)
 |    |-- element: struct (containsNull = false)
 |    |    |-- inicial: integer (nullable = true)
 |    |    |-- final: integer (nullable = true)

df2数据： #

seller_id  cep
12222      114
33332      43
33344      338

期望结果df_final： #

state   interval      limits       cep_ori              seller_id
PA     >=5, <=10      >10    [{111, 222}, {333, 444}]     12222
PA     >=5, <=10      >10    [{111, 222}, {333, 444}]     33344
SP     >=6, <=10      <8     [{333, 444}, {555, 666}]     33344

我尝试了以下代码但未成功：

df.filter(sf.size(sf.array_intersect(df.cep_ori.inicial, sf.array(df2.cep)) != 0))
.filter(sf.size(sf.array_intersect(df.cep_ori.final,sf.array(df2.cep)) != 0))

你的代码思路错误，array_intersect是求数组交集，无法直接判断数值是否落在区间内。这里提供两种高效的实现方式：

方法一：利用高阶函数exists（推荐，无需展开数组） #

直接使用Spark的高阶函数exists遍历cep_ori数组，判断df2.cep是否落在任意一个区间内，代码简洁高效：

import pyspark.sql.functions as sf

df_final = df.join(
    df2,
    sf.exists(
        "cep_ori",
        lambda x: (sf.col("cep") >= x["inicial"]) & (sf.col("cep") <= x["final"])
    ),
    how="inner"
)

# 查看结果
df_final.show(truncate=False)

方法二：展开数组后匹配再去重 #

先将cep_ori数组展开为单条区间记录，匹配后再去重保留原始行：

import pyspark.sql.functions as sf

# 展开cep_ori数组为单独的区间行
df_exploded = df.withColumn("cep_range", sf.explode("cep_ori"))

# 提取区间的上下限字段
df_exploded = df_exploded.withColumn("inicial", sf.col("cep_range.inicial")) \
                         .withColumn("final", sf.col("cep_range.final"))

# 内连接df2，匹配区间条件
df_joined = df_exploded.join(
    df2,
    (sf.col("cep") >= sf.col("inicial")) & (sf.col("cep") <= sf.col("final")),
    how="inner"
)

# 保留原始字段与seller_id，去重避免重复行
df_final = df_joined.select("state", "interval", "limits", "cep_ori", "seller_id") \
                    .dropDuplicates(["state", "interval", "limits", "cep_ori", "seller_id"])

# 查看结果
df_final.show(truncate=False)

两种方法都能得到你期望的结果，其中方法一无需展开数组，性能更优。

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