在 Spark 中重新分组/连接 DataFrame 行

Question

我有一个如下所示的 DataFrame：

scala> data.show
+-----+---+---------+
|label| id| features|
+-----+---+---------+
|  1.0|  1|[1.0,2.0]|
|  0.0|  2|[5.0,6.0]|
|  1.0|  1|[3.0,4.0]|
|  0.0|  2|[7.0,8.0]|
+-----+---+---------+

我想根据 "id" 重新组合特征，这样我可以获得以下内容：

scala> data.show
+---------+---+-----------------+
|    label| id| features        |
+---------+---+-----------------+
|  1.0,1.0|  1|[1.0,2.0,3.0,4.0]|
|  0.0,0.0|  2|[5.0,6.0,7.8,8.0]|
+---------+---+-----------------+

这是我用来生成上述 DataFrame 的代码

val rdd = sc.parallelize(List((1.0, 1, Vectors.dense(1.0, 2.0)), (0.0, 2, Vectors.dense(5.0, 6.0)), (1.0, 1, Vectors.dense(3.0, 4.0)), (0.0, 2, Vectors.dense(7.0, 8.0))))
val data = rdd.toDF("label", "id", "features")

我一直在尝试使用 RDD 和 DataFrames 做不同的事情。到目前为止，最 "promising" 的方法是根据 "id"

进行过滤

data.filter($"id".equalTo(1))

+-----+---+---------+
|label| id| features|
+-----+---+---------+
|  1.0|  1|[1.0,2.0]|
|  1.0|  1|[3.0,4.0]|
+-----+---+---------+

但是我现在有两个瓶颈：

1) 如何自动过滤 "id" 可能具有的所有不同值？

以下会产生错误：

data.select("id").distinct.foreach(x => data.filter($"id".equalTo(x)))

2) 如何将常见的 "features" 方面连接到给定的 "id"。没有尝试太多，因为我仍然停留在 1)

欢迎任何建议

注意：为澄清起见，"label" 对于每次出现的 "id" 始终相同。抱歉造成混淆，我的任务的一个简单扩展也是对 "labels"（更新示例）

进行分组

Answer 1

我认为没有有效的方法来实现您想要的，额外的订单要求并不能使情况变得更好。我能想到的最干净的方法是 groupByKey 像这样：

import org.apache.spark.mllib.linalg.{Vectors, Vector}
import org.apache.spark.sql.functions.monotonicallyIncreasingId
import org.apache.spark.sql.Row
import org.apache.spark.rdd.RDD


val pairs: RDD[((Double, Int), (Long, Vector))] = data
  // Add row identifiers so we can keep desired order
  .withColumn("uid", monotonicallyIncreasingId)
  // Create PairwiseRDD where (label, id) is a key
  // and (row-id, vector is a value)
  .map{case Row(label: Double, id: Int, v: Vector, uid: Long) => 
    ((label, id), (uid, v))}

val rows = pairs.groupByKey.mapValues(xs => {
  val vs = xs
    .toArray
    .sortBy(_._1) // Sort by row id to keep order
    .flatMap(_._2.toDense.values) // flatmap vector values

  Vectors.dense(vs) // return concatenated vectors 

}).map{case ((label, id), v) => (label, id, v)} // Reshape

val grouped = rows.toDF("label", "id", "features")

grouped.show

// +-----+---+-----------------+
// |label| id|         features|
// +-----+---+-----------------+
// |  0.0|  2|[5.0,6.0,7.0,8.0]|
// |  1.0|  1|[1.0,2.0,3.0,4.0]|
// +-----+---+-----------------+

也可以使用类似于我为提议的 UDAF，但它的效率甚至低于这个。

在 Spark 中重新分组/连接 DataFrame 行

Regrouping / Concatenating DataFrame rows in Spark

scala

dataframe

apache-spark

apache-spark-sql

apache-spark-ml