如何在PySpark中调用math.isclose？疑似需广播但不知代码实现

PySpark的列是分布式数据结构，无法直接使用math.isclose或numpy.isclose，但可以通过自定义UDF或内置函数组合实现等价逻辑，针对不同场景有不同实现方式：

1. 逐行比较两列的值 #

如果需要比较DataFrame中两列的对应行是否近似相等，直接将math.isclose包装成UDF即可：

代码示例 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import BooleanType
import math

# 初始化SparkSession
spark = SparkSession.builder.appName("isclose_demo").getOrCreate()

# 创建测试数据
test_data = [(1.0000001, 1.0), (2.5, 2.5000003), (3.0, 3.1)]
df = spark.createDataFrame(test_data, ["col_a", "col_b"])

# 定义UDF，封装math.isclose逻辑，可自定义rel_tol和abs_tol参数
isclose_udf = udf(
    lambda x, y: math.isclose(x, y, rel_tol=1e-09, abs_tol=0.0),
    BooleanType()
)

# 应用UDF生成结果列
df = df.withColumn("values_are_close", isclose_udf(df.col_a, df.col_b))
df.show()

输出结果：

+---------+---------+----------------+
|    col_a|    col_b|values_are_close|
+---------+---------+----------------+
|1.0000001|      1.0|            true|
|      2.5|2.5000003|            true|
|      3.0|      3.1|           false|
+---------+---------+----------------+

2. 比较列与单个常量值（可选广播优化） #

如果需要将列的所有值与一个固定常量比较，可结合广播变量优化（对于简单数值其实无需广播，但如果是复杂大对象，广播能减少数据传输开销）：

代码示例 #

# 广播目标常量值
target_num = spark.sparkContext.broadcast(1.0)

# 定义针对常量的UDF
isclose_to_const_udf = udf(
    lambda x: math.isclose(x, target_num.value, rel_tol=1e-09, abs_tol=0.0),
    BooleanType()
)

# 应用UDF
df = df.withColumn("close_to_1", isclose_to_const_udf(df.col_a))
df.show()

# 使用完毕后销毁广播变量（可选）
target_num.unpersist()

3. 用PySpark内置函数实现等价逻辑 #

如果不想用UDF，也可以手动用内置函数组合出math.isclose的逻辑（同时判断绝对误差和相对误差）：

from pyspark.sql.functions import abs, lit, or_

rel_tol = lit(1e-09)
abs_tol = lit(0.0)

df = df.withColumn(
    "is_close_builtin",
    or_(
        abs(df.col_a - df.col_b) <= abs_tol,
        abs(df.col_a - df.col_b) <= rel_tol * abs(df.col_a),
        abs(df.col_a - df.col_b) <= rel_tol * abs(df.col_b)
    )
)
df.show()

这个写法和UDF的效果完全一致，适合偏好纯内置函数的场景。

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