«ожидаемые нулевые аргументы для построения ClassDict (для numpy.dtype)» при вызове UDF, который возвращает FloatType ()

Question

Я полагаю, что это связано с этим: Ошибка искры: ожидаемые нулевые аргументы для построения ClassDict (для numpy.core.multiarray._reconstruct)

У меня есть фрейм данных

id col_1 col_2
1 [1,2] [1,3]
2 [2,1] [3,4]

Я хочу создать еще один столбец с расстоянием cosine между col_1 и col_2.

from scipy.spatial.distance import cosine

def cosine_distance(a,b):
    try:
        return cosine(a, b)
    except Exception as e:
        return 0.0 # in case division by zero

И я определил udf:

cosine_distance_udf = udf (cosine_distance, FloatType())

И наконец:

new_df = df.withColumn('cosine_distance', cosine_distance_udf('col_1', 'col_2'))

И у меня ошибка: PickleException: expected zero arguments for construction of ClassDict (for numpy.dtype)

Что я сделал не так?

Answer 1

Причина ошибки ясна при проверке типа возвращаемого значения cosine:

type(cosine([1, 2], [1, 3]))
# numpy.float64

Однако np.float64 является подклассом float:

issubclass(np.float64, float)
# True

Итак, с небольшим изменением вашей функции,

def cosine_distance(a, b):
    try:
        return float(cosine(a, b)) # cosine(a, b).item()
    except Exception as e:
        return 0.0 # in case division by zero

Это будет работать

df.withColumn('cosine_distance', cosine_distance_udf('col_1', 'col_2')).show()

+------+------+---------------+
| col_1| col_2|cosine_distance|
+------+------+---------------+
|[1, 2]|[3, 4]|     0.01613009|
|[2, 1]|[3, 4]|     0.10557281|
+------+------+---------------+