在PySpark中为RDD使用UDF遇类型错误的解决求助

代码与错误 #

自定义距离函数及注册UDF #

from pyspark.sql.functions import abs as pyspark_abs, sum as pyspark_sum, \
                    rand as pyspark_rand, min as pyspark_min, max as pyspark_max, \
                    radians as pyspark_radians, cos as pyspark_cos, sin as pyspark_sin, \
                    sqrt as pyspark_sqrt, atan2 as pyspark_atan2

def distance(point1, point2):
    lat1, lon1 = point1
    lat2, lon2 = point2
    R = 6371  # 地球半径，单位千米
    dlat = pyspark_radians(lat2 - lat1)
    dlon = pyspark_radians(lon2 - lon1)
    a = pyspark_sin(dlat / 2) * pyspark_sin(dlat / 2) + pyspark_cos(pyspark_radians(lat1)) * pyspark_cos(pyspark_radians(lat2)) * pyspark_sin(dlon / 2) * pyspark_sin(dlon / 2)
    c = 2 * pyspark_atan2(pyspark_sqrt(a), pyspark_sqrt(1 - a))
    distance = R * c
    return distance

spark.udf.register("distance", distance)

RDD调用代码 #

neighbors = radian_rdd.cartesian(radian_rdd) \
                   .filter(lambda x: x[0][1] < x[1][1]) \
                   .filter(lambda x: distance(x[0][0], x[1][0]) <= eps) \
                   .map(lambda x: (x[0][1], x[1][1]))

错误信息 #

PySparkTypeError(pyspark.errors.exceptions.base.PySparkTypeError: [NOT_COLUMN_OR_STR] Argument `col` should be a Column or str, got float.)

问题根源 #

你导入的pyspark.sql.functions下的函数（如pyspark_radians、pyspark_sin等）是专门用于操作DataFrame的Column对象的API，而RDD中的元素是原生Python类型（比如float），直接传入这些SQL函数会触发类型不匹配错误。另外，spark.udf.register是将函数注册为DataFrame/SQL可用的UDF，对RDD操作完全没用，属于多余步骤。

修复方案 #

改用Python标准库math模块中的对应函数来实现距离计算，因为RDD操作基于原生Python对象，适配math模块的函数：

修改后的代码 #

import math

def distance(point1, point2):
    lat1, lon1 = point1
    lat2, lon2 = point2
    R = 6371  # 地球半径，单位千米
    dlat = math.radians(lat2 - lat1)
    dlon = math.radians(lon2 - lon1)
    a = math.sin(dlat / 2) ** 2 + math.cos(math.radians(lat1)) * math.cos(math.radians(lat2)) * math.sin(dlon / 2) ** 2
    c = 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))
    return R * c

# RDD调用代码无需修改，直接使用上述函数
neighbors = radian_rdd.cartesian(radian_rdd) \
                   .filter(lambda x: x[0][1] < x[1][1]) \
                   .filter(lambda x: distance(x[0][0], x[1][0]) <= eps) \
                   .map(lambda x: (x[0][1], x[1][1]))

说明 #

• 移除了所有pyspark.sql.functions的导入，替换为math模块的函数，这些函数接受原生数值类型（如float）作为参数，完全适配RDD的元素类型。
• 删除了spark.udf.register语句，因为RDD不需要注册UDF，直接调用Python函数即可。

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