Перебирать элементы столбцов Scala

Question

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

 A      B 
(1,2)  (1,3)
(2,3)  (3,4)

Создать:

 A      B     C
(1,2)  (1,3)  1
(2,3)  (3,4)  1.4

Моя схема данных:

df.schema.foreach(println)
StructField(col1,ArrayType(DoubleType,false),false)
StructField(col2,ArrayType(DoubleType,false),true)

Всякий раз, когда я вызываю эту функцию расстояния:

def distance(xs: Array[Double], ys: Array[Double]) = {
  sqrt((xs zip ys).map { case (x,y) => pow(y - x, 2) }.sum)
}

Я получаю ошибку типа:

df.withColumn("distances" , distance($"col1",$"col2"))
<console>:68: error: type mismatch;
 found   : org.apache.spark.sql.ColumnName
 required: Array[Double]
       ids_with_predictions_centroids3.withColumn("distances" , distance($"col1",$"col2"))

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

Answer 1

Функции Spark работают на основе столбцов и Ваша единственная ошибка в том, что вы смешиваете столбец и примитивы в функции

И сообщение об ошибке достаточно четкое, в котором говорится, что вы передаете столбец в функции distance , т.е. $"col1" и $"col2" - это column , но расстояние функция определяется как distance(xs: Array[Double], ys: Array[Double]), принимающая примитивных типов .

Решение состоит в том, чтобы сделать функцию расстояния полностью основанной на столбце как

import org.apache.spark.sql.Column
import org.apache.spark.sql.functions._

def distance(xs: Column, ys: Column) = {
  sqrt(pow(ys(0)-xs(0), 2) + pow(ys(1)-xs(1), 2))
}

df.withColumn("distances" , distance($"col1",$"col2")).show(false)

, который должен дать вам правильный результат без ошибок

+------+------+------------------+
|col1  |col2  |distances         |
+------+------+------------------+
|[1, 2]|[1, 3]|1.0               |
|[2, 3]|[3, 4]|1.4142135623730951|
+------+------+------------------+

Надеюсь, ответ полезен

Answer 2

Чтобы использовать пользовательскую функцию в кадре данных, необходимо определить ее как UDF. Это можно сделать, например, следующим образом:

val distance = udf((xs: WrappedArray[Double], ys: WrappedArray[Double]) => {
  math.sqrt((xs zip ys).map { case (x,y) => math.pow(y - x, 2) }.sum)
})

df.withColumn("C", distance($"A", $"B")).show()

Обратите внимание, что WrappedArray (или Seq) необходимо использовать здесь.

Результирующий кадр данных:

+----------+----------+------------------+
|         A|         B|                 C|
+----------+----------+------------------+
|[1.0, 2.0]|[1.0, 3.0]|               1.0|
|[2.0, 3.0]|[3.0, 4.0]|1.4142135623730951|
+----------+----------+------------------+