Преобразование произвольного количества столбцов в вектор

77
8

Как преобразовать группу произвольных столбцов в вектор Mllib? По сути, у меня есть первый столбец моего DataFrame с фиксированным именем, а затем количество произвольных именованных столбцов, каждый из которых имеет значения Double внутри.

вот так:

name  |  a  |  b  |  c  |
val1 | 0.0 | 1.0 | 1.0 |
val2 | 2.0 | 1.0 | 5.0 |

Может быть любое количество столбцов. Мне нужно получить DataSet следующего:

final case class ValuesRow(name: String, values: Vector)

спросил(а) 2021-01-25T17:54:11+03:00 4 месяца, 4 недели назад
1
Решение
77

Это можно сделать простым способом с помощью VectorAssembler. В качестве входных данных используются столбцы, которые должны быть объединены в Vector, в этом случае все столбцы, кроме первого.

val df = spark.createDataFrame(Seq(("val1", 0, 1, 1), ("val2", 2, 1, 5)))
.toDF("name", "a", "b", "c")

val columnNames = df.columns.drop(1) // drop the name column
val assembler = new VectorAssembler()
.setInputCols(columnNames)
.setOutputCol("values")

val df2 = assembler.transform(df).select("name", "values").as[ValuesRow]

Результатом будет набор данных, содержащий столбцы имен и значений:

+----+-------------+
|name| values|
+----+-------------+
|val1|[0.0,1.0,1.0]|
|val2|[2.0,1.0,5.0]|
+----+-------------+

ответил(а) 2021-01-25T17:54:11+03:00 4 месяца, 4 недели назад
63

Вот один из способов сделать это:

import org.apache.spark.sql.functions._
import org.apache.spark.mllib.linalg.DenseVector

val ds = Seq(
("val1", 0.0, 1.0, 1.0),
("val2", 2.0, 1.0, 5.0)
).toDF("name", "a", "b", "c").
as[(String, Double, Double, Double)]

val colList = ds.columns
val keyCol = colList(0)
val valCols = colList.drop(1)

def arrToVec = udf(
(s: Seq[Double]) => new DenseVector(s.toArray)
)

ds.select(
col(keyCol), arrToVec( array(valCols.map(x => col(x)): _*) ).as("values")
).show
// +----+-------------+
// |name| values|
// +----+-------------+
// |val1|[0.0,1.0,1.0]|
// |val2|[2.0,1.0,5.0]|
// +----+-------------+

ответил(а) 2021-01-25T17:54:11+03:00 4 месяца, 4 недели назад
Ваш ответ
Введите минимум 50 символов
Чтобы , пожалуйста,
Выберите тему жалобы:

Другая проблема