Spark Dataframe - получить все списки пар (Scala)

Question

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

Это исходный фрейм данных:

+---------------+
|  candidateList|
+---------------+
|         [1, 2]|
|      [2, 3, 4]|
|      [1, 3, 5]|
|[1, 2, 3, 4, 5]|
|[1, 2, 3, 4, 5]|
+---------------+

А вот так должно выглядеть после вычисления:

+---------------+
|  candidates   |
+---------------+
|         [1, 2]|
|         [2, 3]|
|         [2, 4]|
|         [3, 4]|
|         [1, 3]|
|         [1, 5]|
|         [3, 5]|
|and so on...   |
+---------------+

Я действительно не знаю, как это возможно в искре, может, у кого-то есть совет для меня.

С уважением

Седир Мухаммед

Answer 1

Вам нужно создать UDF (пользовательскую функцию) и использовать ее с функцией explode. Сам UDF прост благодаря методу combinations коллекции Scala:

import scala.collection.mutable
import org.apache.spark.sql.functions._
import spark.implicits._

val pairsUdf = udf((arr: mutable.Seq[Int]) => arr.combinations(2).toArray)
val result = df.select(explode(pairsUdf($"candidateList")) as "candidates")

result.show(numRows = 8)
// +----------+
// |candidates|
// +----------+
// |    [1, 2]|
// |    [2, 3]|
// |    [2, 4]|
// |    [3, 4]|
// |    [1, 3]|
// |    [1, 5]|
// |    [3, 5]|
// |    [1, 2]|
// +----------+