Умножение двух столбцов в фрейме данных pyspark.Один из них содержит вектор, а один из них содержит константу

Question

У меня есть фрейм данных pyspark, в котором есть один столбец с векторными значениями и один столбец с постоянными числовыми значениями.Например,

A | B
1 | [2,4,5]
5 | [6,5,3] 

Я хочу умножить столбец вектора на столбец константы.Я пытаюсь сделать это в основном потому, что у меня есть слова wmbeddings в столбце B и некоторые веса в столбце A.И моя последняя цель - получить взвешенные вложения.

Answer 1

Spark 2,4 и далее, вы можете использовать функции высшего порядка, доступные в sql.

scala> val df = Seq((1,Seq(2,4,5)),(5,Seq(6,5,3))).toDF("a","b")
df: org.apache.spark.sql.DataFrame = [a: int, b: array<int>]

scala> df.createOrReplaceTempView("ashima")

scala> spark.sql(""" select a, b, transform(b, x -> x * a) as result from ashima """).show(false)
+---+---------+------------+
|a  |b        |result      |
+---+---------+------------+
|1  |[2, 4, 5]|[2, 4, 5]   |
|5  |[6, 5, 3]|[30, 25, 15]|
+---+---------+------------+


scala>

Answer 2

Если ваши векторные данные хранятся в виде массива значений типа double, вы можете сделать это:

import breeze.linalg.{Vector => BV}

val data = spark.createDataset(Seq(
    (1, Array[Double](2, 4, 5)),
    (5, Array[Double](6, 5, 3))
  )).toDF("A", "B")

data.as[(Long, Array[Double])].map(r => {
  (BV(r._2) * r._1.toDouble).toArray
}).show()

, который становится

+------------------+
|             value|
+------------------+
|   [2.0, 4.0, 5.0]|
|[30.0, 25.0, 15.0]|
+------------------+