来自 Array 列的所有组合的 Spark Dataframe

Question

假设我有一个 Spark DataFrame d1 有两列 elements_1 和 elements_2，其中包含大小为 k 的整数集，以及包含整数值的 value_1、value_2。例如，k = 3:

d1 = 
+------------+------------+
| elements_1 | elements_2 |
+-------------------------+
| (1, 4, 3)  |  (3, 4, 5) |
| (2, 1, 3)  |  (1, 0, 2) |
| (4, 3, 1)  |  (3, 5, 6) |
+-------------------------+

我需要创建一个新列combinations，其中包含对于每对集合elements_1 和elements_2，来自其元素的所有可能组合的集合列表。这些集合必须具有以下属性：

1. 它们的大小必须是k+1
2. 它们必须包含elements_1 中的集合或elements_2 中的集合

例如，从(1, 2, 3) 和(3, 4, 5) 我们得到[(1, 2, 3, 4), (1, 2, 3, 5), (3, 4, 5, 1) and (3, 4, 5, 2)]。该列表不包含(1, 2, 5)，因为它的长度不是3+1，也不包含(1, 2, 4, 5)，因为它既不包含原始集合。

Answer 1

您需要创建一个自定义的用户定义函数来执行转换，从中创建一个与 spark 兼容的UserDefinedFunction，然后使用 withColumn 应用。所以真的，这里有两个问题：（1）如何进行你描述的集合转换，以及（2）如何使用用户定义的函数在 DataFrame 中创建一个新列。

这是设置逻辑的第一个镜头，如果它符合您的要求，请告诉我：

def combo[A](a: Set[A], b: Set[A]): Set[Set[A]] = 
    a.diff(b).map(b+_) ++ b.diff(a).map(a+_)

现在创建 UDF 包装器。请注意，在后台这些集合都由 WrappedArrays 表示，因此我们需要处理这个问题。通过定义一些隐式转换，可能有一种更优雅的方式来处理这个问题，但这应该可以：

import scala.collection.mutable.WrappedArray
val comboWrap: (WrappedArray[Int],WrappedArray[Int])=>Array[Array[Int]] = 
    (x,y) => combo(x.toSet,y.toSet).map(_.toArray).toArray
val comboUDF = udf(comboWrap)

最后，通过创建一个新列将其应用到 DataFrame：

val data = Seq((Set(1,2,3),Set(3,4,5))).toDF("elements_1","elements_2")
val result = data.withColumn("result", 
    comboUDF(col("elements_1"),col("elements_2")))
result.show