Преобразовать столбец искровой строки в vectorUDT

Question

Я использую pyspark.
Сначала я читаю CSV со строковым столбцом. Данные внутри выглядят так:

(174, [7, 10, 56, 89, 156], [1.0, 1.0, 1.0, 1.0, 1.0])

Мне нужно преобразовать его в VectorUDT, чтобы дать столбец в качестве входных данных для алгоритмов машинного обучения.

Я уже пытался привести столбец:

data = data.withColumn (column, data [column] .cast (VectorUDT ())

Но это не работает ...

У вас есть какие-нибудь решения?

Answer 1

Один из способов, который вы можете попробовать, - это использовать внутренний _sqlType (см. Ниже) VectorUDT для создания строки JSON, а затем использовать функцию from_json:

struct<type:tinyint,size:int,indices:array<int>,values:array<double>> 

Сначала удалите начальные и конечные скобки и разделите строку с помощью регулярного выражения ,\s*(?=\[). это преобразует столбец StringType в столбец ArrayType с именем s с 3 элементами, соответствующими s [0] = размер , s [1] = индексы и s [ 2] = значения , а затем используйте функцию concat для создания строки JSON.

Примечание: в этом примере мы заключаем VectorUDT с массив , поскольку функция from_ json принимает только один из сложных типов данных: массив , map или struct . Вы можете попробовать использовать map или struct в качестве оболочки.

from pyspark.ml.linalg import VectorUDT
from pyspark.sql.types import ArrayType
from pyspark.sql.functions import expr, from_json

df = spark.createDataFrame([('(174, [7, 10, 56, 89, 156], [1.0, 1.0, 1.0, 1.0, 1.0])',)],['column'])
# DataFrame[column: string]

df_new = df.withColumn("s", expr("split(substr(column,2,length(column)-2), ',\\\\s*(?=\\\\[)')")) \
  .selectExpr("""
      concat(
        /* type = 0 for SparseVector and type = 1 for DenseVector */
        '[{"type":0,"size":',
        s[0],
        ',"indices":',
        s[1],
        ',"values":',
        s[2],
        '}]'
      ) as vec_json
   """) \
  .withColumn('features', from_json('vec_json', ArrayType(VectorUDT()))[0])

Обратите внимание , что строка JSON, которую мы здесь создали массив с одним VectorUDT, мы можем использовать from_json и getItem(0) для извлечения вектора.

result:

df_new.printSchema()
root
 |-- vec_json: string (nullable = true)
 |-- features: vector (nullable = true)

df_new.show(truncate=False, vertical=True)
-RECORD 0---------------------------------------------------------------------------------------------
 vec_json | [{"type":0,"size":174,"indices":[7, 10, 56, 89, 156],"values":[1.0, 1.0, 1.0, 1.0, 1.0]}]
 features | (174,[7,10,56,89,156],[1.0,1.0,1.0,1.0,1.0])