将多列映射到 Spark 数据框中的单个键

Question

我有一个如下所示的 Spark 数据框：

+------+-----+-----+
|acctId|vehId|count|
+------+-----+-----+
|     1|  666|    1|
|     1|  777|    3|
|     1|  888|    2|
|     1|  999|    3|
|     2|  777|    1|
|     2|  888|    3|
|     2|  999|    1|
|     3|  777|    4|
|     3|  888|    2|
+------+-----+-----+

我想将每个 acctId 的 vehId 映射到其计数，并将其存储回数据框中，因此最终结果如下所示：

+------+---------------------------------------------+
|acctId| map                                         |
+------+---------------------------------------------+
|     1| Map(666 -> 1, 777 -> 3, 888 -> 2, 999 -> 3) |
|     2| Map(777 -> 1, 888 -> 3, 999 -> 1)           |
|     3| Map(777 -> 4, 888 -> 2)                     |
+------+---------------------------------------------+

最好的方法是什么？

我尝试将数据框转换为 RDD 并在行上执行映射，但我不确定如何将每个映射聚合回单个 acctId。总的来说，我是 Spark 和数据帧的新手，但我已尽我所能尝试找到类似的问题——如果这是一个非常常见的问题，我深表歉意。

供您参考/使用，这是我生成测试数据的方式：

val testData = Seq(
    (1, 999),
    (1, 999),
    (2, 999),
    (1, 888),
    (2, 888),
    (3, 888),
    (2, 888),
    (2, 888),
    (1, 888),
    (1, 777),
    (1, 666),
    (3, 888),
    (1, 777),
    (3, 777),
    (2, 777),
    (3, 777),
    (3, 777),
    (1, 999),
    (3, 777),
    (1, 777)
).toDF("acctId", "vehId")

val grouped = testData.groupBy("acctId", "vehId").count

Answer 1

我认为您必须使用 double groupBy 如下所示

val testData = Seq(
  (1, 999),
  (1, 999),
  (2, 999),
  (1, 888),
  (2, 888),
  (3, 888),
  (2, 888),
  (2, 888),
  (1, 888),
  (1, 777),
  (1, 666),
  (3, 888),
  (1, 777),
  (3, 777),
  (2, 777),
  (3, 777),
  (3, 777),
  (1, 999),
  (3, 777),
  (1, 777)
).toDF("acctId", "vehId")

//udf to convert list to map
val listToMap = udf((input: Seq[Row]) => input.map(row => (row.getAs[Int](0), row.getAs[Long](1))).toMap)

val resultDF = testData.groupBy("acctId", "vehId")
  .agg(count("acctId").cast("long").as("count"))
  .groupBy("acctId")
  .agg(collect_list(struct("vehId", "count")) as ("map"))
  .withColumn("map", listToMap($"map"))

输出：

resultDF.show(false)
+------+----------------------------------------+
|acctId|map                                     |
+------+----------------------------------------+
|1     |[777 -> 3, 666 -> 1, 999 -> 3, 888 -> 2]|
|3     |[777 -> 4, 888 -> 2]                    |
|2     |[777 -> 1, 999 -> 1, 888 -> 3]          |
+------+----------------------------------------+

架构：

resultDF.printSchema()
root
 |-- acctId: integer (nullable = false)
 |-- map: map (nullable = true)
 |    |-- key: integer
 |    |-- value: long (valueContainsNull = false)