Spark Scala: Как обновить каждый столбец DataFrame в соответствии с каждой позицией вектора

Question

У меня есть DF, как это:

+--------------------+-----+--------------------+
|               col_0|col_1|               col_2|
+--------------------+-----+--------------------+
|0.009069428120139292|  0.3|9.015488712438252E-6|
|0.008070826019024355|  0.4|3.379696051366339...|
|0.009774715414895803|  0.1|1.299590589291292...|
|0.009631155146285946|  0.9|1.218569739510422...|

и два вектора:

v1[7.0,0.007,0.052]
v2[804.0,553.0,143993.0]

Общее количество столбцов совпадает с общим количеством позиций в каждом векторе. Как можно применить уравнение, используя числа, сохраненные в i-й позиции, чтобы сделать некоторые вычисления для обновления текущего значения DF (в i-й позиции)? Я имею в виду, мне нужно обновить все значения в DF, используя значения в векторах.

Answer 1

Возможно, что-то вроде этого - то, что вы ищете?

import org.apache.spark.sql.Column
import org.apache.spark.sql.DataFrame

val df = Seq((1,2,3),(4,5,6)).toDF

val updateVector = Vector(10,20,30)

val updateFunction = (columnValue: Column, vectorValue: Int) => columnValue * lit(vectorValue)

val updateColumns = (df: DataFrame, updateVector: Vector[Int], updateFunction:((Column, Int) => Column)) => {
    val columns = df.columns
    updateVector.zipWithIndex.map{case (updateValue, index) => updateFunction(col(columns(index)), updateVector(index)).as(columns(index))}
}

val dfUpdated = df.select(updateColumns(df, updateVector, updateFunction) :_*)

dfUpdated.show

+---+---+---+
| _1| _2| _3|
+---+---+---+
| 10| 40| 90|
| 40|100|180|
+---+---+---+