Конвертировать PySpark DenseVector в массив

Question

Я пытаюсь преобразовать столбец данных pyspark в DenseVector в массив, но всегда получаю сообщение об ошибке.

data = [(Vectors.dense([8.0, 1.0, 3.0, 2.0, 5.0]),),
(Vectors.dense([2.0, 0.0, 3.0, 4.0, 5.0]),),
(Vectors.dense([4.0, 0.0, 0.0, 6.0, 7.0]),)]

df = spark.createDataFrame(data,["features"])

Я пытался определить UDF и использовать toArray ()

to_array = udf(lambda x: x.toArray(), ArrayType(FloatType()))
df = df.withColumn('features', to_array('features'))

Но потом я получаю следующую ошибку, если я выполняю df.collect ()

org.apache.spark.SparkException: Job aborted due to stage failure: Task 1 in stage 17.0 failed 4 times, 
most recent failure: Lost task 1.3 in stage 17.0 (TID 100, 10.139.64.6, executor 0): 
net.razorvine.pickle.PickleException: expected zero arguments for construction of ClassDict 
(for numpy.core.multiarray._reconstruct)

Есть идеи, как мне этого добиться?

Answer 1

toArray () возвращает numpy.ndarray, который неявным образом не может быть преобразован в ArrayType(FloatType()). Используйте дополнительно .tolist() для преобразования:

#or: to_array = F.udf(lambda v: list([float(x) for x in v]), T.ArrayType(T.FloatType()))
to_array = F.udf(lambda v: v.toArray().tolist(), T.ArrayType(T.FloatType()))
df = df.withColumn('features', to_array('features'))