scala - 如何获取 Spark Dataframe 上按结果分组的元组？

Question

我正在尝试根据 id 对实体进行分组，运行以下代码我有这个数据框:

val pet_type_count = pet_list.groupBy("id","pets_type").count()
pet_type_count.sort("id").limit(20).show

+----------+---------------------+-----+
|        id|            pets_type|count|
+----------+---------------------+-----+
|         0|                    0|    2|
|         1|                    0|    3|
|         1|                    3|    3|
|        10|                    0|    4|
|        10|                    1|    1|
|        13|                    0|    3|
|        16|                    1|    3|
|        17|                    1|    1|
|        18|                    1|    2|
|        18|                    0|    1|
|        19|                    1|    7|
+----------+---------------------+-----+

我想按 id 对组的结果进行分组，现在返回每个 id 的元组列表，这样我就可以为每个 id 应用以下 udf:

val agg_udf =  udf { (v1: List[Tuple2[String, String]]) =>
    var feature_vector = Array.fill(5)(0)
    for (row <- v1) {
      val index = (5 - row._1.toInt)
      vector(index) = row._2.toInt
    }
    vector
}

val pet_vector_included = pet_type_count.groupBy("id").agg(agg_udf(col("pets_type_count")).alias("pet_count_vector"))

为此我需要获得以下内容:

+----------+---------------------+-----+
|        id|            pets_type_count|
+----------+---------------------+-----+
|         0|                      (0,2)|
|         1|                      (0,3)|
|          |                      (3,3)|
|        10|                      (0,4)|
|          |                      (1,1)|
|        13|                      (0,3)|
|        16|                      (1,3)|
|        17|                      (1,1)|
|        18|                      (1,2)|
|          |                      (0,1)|
|        19|                      (1,7)|
+----------+---------------------+-----+

我无法弄清楚如何在 id 上的 groupby 之后获取元组。任何帮助将不胜感激!

Answer 1

您可以简单地使用 struct 内置函数 将 pets_type 和 count 列作为一列并使用 collect_list 内置函数 用于在按 id 分组时收集新形成的列。您可以 orderBy 仅按 id 列对 dataframe 进行排序。

import org.apache.spark.sql.functions._
val pet_type_count = df.withColumn("struct", struct("pets_type", "count"))
  .groupBy("id").agg(collect_list(col("struct")).as("pets_type_count"))
  .orderBy("id")

这应该会给你你想要的结果

+---+---------------+
|id |pets_type_count|
+---+---------------+
|0  |[[0,2]]        |
|1  |[[0,3], [3,3]] |
|10 |[[0,4], [1,1]] |
|13 |[[0,3]]        |
|16 |[[1,3]]        |
|17 |[[1,1]]        |
|18 |[[1,2], [0,1]] |
|19 |[[1,7]]        |
+---+---------------+

所以你可以应用你定义的udf函数(它也需要一些修改)如下

val agg_udf =  udf { (v1: Seq[Row]) =>
  var feature_vector = Array.fill(5)(0)
  for (row <- v1) {
    val index = (4 - row.getAs[Int](0))
    feature_vector(index) = row.getAs[Int](1)
  }
  feature_vector
}

val pet_vector_included = pet_type_count.withColumn("pet_count_vector", agg_udf(col("pets_type_count")))

pet_vector_included.show(false)

应该给你

+---+---------------+----------------+
|id |pets_type_count|pet_count_vector|
+---+---------------+----------------+
|0  |[[0,2]]        |[0, 0, 0, 0, 2] |
|1  |[[0,3], [3,3]] |[0, 3, 0, 0, 3] |
|10 |[[0,4], [1,1]] |[0, 0, 0, 1, 4] |
|13 |[[0,3]]        |[0, 0, 0, 0, 3] |
|16 |[[1,3]]        |[0, 0, 0, 3, 0] |
|17 |[[1,1]]        |[0, 0, 0, 1, 0] |
|18 |[[1,2], [0,1]] |[0, 0, 0, 2, 1] |
|19 |[[1,7]]        |[0, 0, 0, 7, 0] |
+---+---------------+----------------+

希望回答对你有帮助