数组元素组合

Question

我有一个array[Int]的rdd，像下面这样。

([0,1,7],[0,1],[0,1,3],...)

现在我想得到每个内表中的数组组合，就像下面这样。

Array [[0,1,7], [0,7],[1,7],[7],[0,1],[1],[0,1,3],[0,3],[1,3],[3]]

怎样才能做到这一点？

Answer 1

你可以使用 combinations 从 Seq 特质和 flatMap 从Dataset API中输出。

import spark.implicits._
val df = Seq(
  Array(0, 1, 7),
  Array(0, 1),
  Array(0, 1, 3)
).toDF("arr")
df.show()
val res = df.flatMap{
  row =>
    val rowSeq: Seq[Int] = row.getAs[Seq[Int]](0)
    (1 to rowSeq.length).flatMap(n => rowSeq.combinations(n))
}.collect().foldLeft(Set.empty[Array[Int]]){
 case (arrAcc, seqInt) => arrAcc + seqInt.toArray
}
println(res.map(a => a.mkString("[", ", ", "]")).mkString("[", ", ", "]"))

输出将是：

df.show()
+---------+
|      arr|
+---------+
|[0, 1, 7]|
|   [0, 1]|
|[0, 1, 3]|
+---------+
res
[[1], [1, 7], [0], [0, 7], [0, 1], [3], [0, 1, 3], [0, 3], [7], [0, 1, 7], [1, 3]]