Как преобразовать строковый массив в массив numpy и передать его в UDF в Pyspark?

Question

Я сохранил массив Numpy в виде строкового массива в CSV-файле (другого способа не знал). Теперь есть две проблемы, с которыми я сталкиваюсь

1) Мне нужно прочитать файл CSV, преобразовать массив строк в массив numpy и передать его в UDF.

2) Почему я не могу использовать метод DF.withcolumn? Выдает сообщение об ошибке ниже

Traceback (последний вызов был последним): файл "", строка 1, в файле "/usr/lib/spark/python/pyspark/sql/dataframe.py", строка 1989, в withColumn assert isinstance (col, Column), «col должно быть Column» AssertionError: col должно быть Column

Мой фрагмент кода -

def wantNumpyArr(array):
    try:
      //some code//
    except Exception:
       pass
    else:
      return float_var

spark.udf.register("wantNumpyArr", wantNumpyArr, FloatType())

#Read from csv file
read_data=spark.read.format("csv").load("/path/to/file/part-*.csv", header="true")
rdd = read_data.rdd

convert_data = rdd.map(lambda x: (x[0], x[1], wantNumpyArr(x[2])))

Когда я печатаю convert_data RDD всегда имеет значение 3-го столбца как None, что означает, что поток в UDF всегда идет в блоке Except. Пример данных -

[Row(Id='ABCD505936', some_string='XCOYNZGAE', array='[0, 2, 5, 6, 8, 10, 12, 13, 14, 15]')]

Схема DF - -

print (read_data.schema)
StructType(List(StructField(col1,StringType,true),StructField(col2,StringType,true),StructField(col3,StringType,true)))

Answer 1

Вы можете легко преобразовать строку обратно в массив, убрав квадратные скобки и разделив их запятыми:

df = spark.createDataFrame([("ABCD505936", "XCOYNZGAE", "[0, 2, 5, 6, 8, 10, 12, 13, 14, 15]")], ['Id', 'some_string', 'array'])

df = df.withColumn("array", split(regexp_replace(col("array"), '[\\[\\]]', ''), ',').cast("array<float>"))
df.printSchema()

#root
# |-- Id: string (nullable = true)
# |-- some_string: string (nullable = true)
# |-- array: array (nullable = true)
# |    |-- element: float (containsNull = true)