Как разбить столбец векторов на два столбца?

Question

Я использую PySpark.

Выход случайного леса Spark ML DataFrame имеет столбец «вероятность», который представляет собой вектор с двумя значениями. Я просто хочу добавить два столбца в выходной DataFrame, «prob1» и «prob2», которые соответствуют первому и второму значениям вектора.

Я попытался следующие:

output2 = output.withColumn('prob1', output.map(lambda r: r['probability'][0])) 

, но я получаю ошибку, что «цв должен быть столбец».

Любые предложения о том, как преобразовать столбец векторов в столбцы его значений?

Answer 1

Возможно, вы захотите использовать один UDF для извлечения первого значения, а другой - для извлечения второго. Затем вы можете использовать UDF с выбранным вызовом на выходе случайного кадра данных forrest. Пример:

from pyspark.sql.functions import udf, col 

split1_udf = udf(lambda value: value[0], FloatType()) 
split2_udf = udf(lambda value: value[1], FloatType()) 
output2 = randomForrestOutput.select(split1_udf(col("probability")).alias("c1"), 
            split2_udf(col("probability")).alias("c2")) 

Это должно дать вам dataframe output2, который имеет столбцы c1 и c2, соответствующие первые и вторых значения в списке, хранящемся в вероятности столбца.

Answer 2

Я понял проблему с предложением выше. В pyspark «плотные векторы просто представлены как объекты массива NumPy», поэтому проблема связана с типами python и numpy. Необходимо добавить .item(), чтобы наложить numpy.float64 на float python.

Следующий код работает:

split1_udf = udf(lambda value: value[0].item(), FloatType()) 
split2_udf = udf(lambda value: value[1].item(), FloatType()) 

output2 = randomforestoutput.select(split1_udf('probability').alias('c1'), split2_udf('probability').alias('c2')) 

Или добавить эти столбцы в исходной dataframe:

randomforestoutput.withColumn('c1', split1_udf('probability')).withColumn('c2', split2_udf('probability')) 

Answer 3

получил ту же проблему, ниже код с поправкой на ситуацию, когда у вас есть п вектор длины.

splits = [udf(lambda value: value[i].item(), FloatType()) for i in range(n)] 
out = tstDF.select(*[s('features').alias("Column"+str(i)) for i, s in enumerate(splits)])