如何基于键为元组的区间字典为PySpark DataFrame添加列？

我有如下Python字典：

ranges = {
    (0, 10): '0 - 10',
    (10, 100): '10 - 100',
    (100, float('inf')): '100+'
}

以及如下PySpark DataFrame：

我希望添加一列Range，当Value列的值落在字典键的左闭右开区间时，返回对应的字典值，最终DataFrame应如下所示：

我知道可以用withColumn和when实现，示例如下：

df.withColumn(
    'Range',
     when((col('Value') >= lit(0)) & (col('Value') < lit(10)), '0 - 10')
)

但当字典有成百上千个键值对或需要调整区间数值时，这种方法效率低下，希望得到更优的实现方案。

方法1：动态构建when条件链 #

通过遍历字典键值对自动拼接when条件，无需手动逐个编写，同时保证区间判断顺序正确：

from pyspark.sql import functions as F

# 按区间左边界排序，确保判断逻辑顺序正确
sorted_ranges = sorted(ranges.items(), key=lambda x: x[0][0])

# 初始化条件表达式
range_expr = F.lit(None)
for (lower, upper), label in sorted_ranges:
    if upper == float('inf'):
        # 处理无穷大区间的特殊情况
        condition = F.col('Value') >= F.lit(lower)
    else:
        condition = (F.col('Value') >= F.lit(lower)) & (F.col('Value') < F.lit(upper))
    range_expr = F.when(condition, F.lit(label)).otherwise(range_expr)

# 添加Range列
result_df = df.withColumn('Range', range_expr)
result_df.show()

方法2：广播变量+UDF（适配复杂逻辑场景） #

如果区间匹配逻辑更复杂，或需要复用区间规则，可通过广播变量传递字典，结合UDF实现：

from pyspark.sql import functions as F
from pyspark.sql.types import StringType

# 广播区间字典，减少节点间数据传输
broadcast_ranges = spark.sparkContext.broadcast(ranges)

# 定义匹配区间的UDF
def get_range(value):
    for (lower, upper), label in broadcast_ranges.value.items():
        if lower <= value < upper:
            return label
    return None

range_udf = F.udf(get_range, StringType())

# 添加Range列
result_df = df.withColumn('Range', range_udf(F.col('Value')))
result_df.show()

方案对比 #

• 方法1基于Spark原生表达式，能被Catalyst优化器处理，大数据量下性能更优，是优先推荐方案。
• 方法2更灵活，支持自定义复杂匹配逻辑，但UDF无法被Spark优化，小数据量或复杂场景下适用。

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