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

Posted 2023-04-13

【中文标题】如何在 Spark Dataframe 上获取按结果分组的元组？

【英文标题】：How to get a Tuple for the grouped by result on a Spark Dataframe?

【发布时间】：2018-03-13 04:29:16

【问题描述】：

我正在尝试根据 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 之后获取元组。任何帮助将不胜感激！

【参考方案1】：

您可以简单地使用struct 内置函数 将pets_type 和count 列为一列，然后使用collect_list 内置函数 收集新的按id 分组时形成的列。您可以通过orderBy 来按id 列订购数据框。

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] |
+---+---------------+----------------+

希望回答对你有帮助