Набор данных Spark: возвращает HashMap значений, имеющих одинаковый ключ

+------+-----+
|userID|entID|
+------+-----+
|     0|    5|
|     0|   15|
|     1|    7|
|     1|    3|
|     2|    3|
|     2|    4|
|     2|    5|
|     2|    9|
|     3|   25|
+------+-----+

Я хочу получить результат как {0->(5,15), 1->(7,3),..}

Любая помощь будет оценена.

2 ответа

Решение

Вот ваш стол снова:

  val df = Seq(
    (0,  5),
    (0, 15),
    (1,  7),
    (1,  3),
    (2,  3),
    (2,  4),
    (2,  5),
    (2,  9),
    (3, 25)
  ).toDF("userId", "entId")
  df.show()

Выходы:

+------+-----+
|userId|entId|
+------+-----+
|     0|    5|
|     0|   15|
|     1|    7|
|     1|    3|
|     2|    3|
|     2|    4|
|     2|    5|
|     2|    9|
|     3|   25|
+------+-----+

Теперь вы можете группировать по userId а потом собирать endId в списки, совмещая получившийся столбец со списками как entIds:

  import org.apache.spark.sql.functions._
  val entIdsForUserId = df.
    groupBy($"userId").                        
    agg(collect_list($"entId").alias("entIds"))

  entIdsForUserId.show()

Выход:

+------+------------+
|userId|      entIds|
+------+------------+
|     1|      [7, 3]|
|     3|        [25]|
|     2|[3, 4, 5, 9]|
|     0|     [5, 15]|
+------+------------+

Заказ после groupBy не указано В зависимости от того, что вы хотите с ним сделать, вы можете дополнительно отсортировать его.

Вы можете собрать его в одну карту на главном узле:

  val m = entIdsForUserId.
    map(r => (r.getAs[Int](0), r.getAs[Seq[Int]](1))).
    collect.toMap

это даст вам:

Map(1 -> List(7, 3), 3 -> List(25), 2 -> List(3, 4, 5, 9), 0 -> List(5, 15))

Одним из подходов будет преобразование набора данных в СДР и выполнение groupByKey, Чтобы получить результат в виде Mapвам нужно будет collect сгруппированный RDD, если набор данных не слишком большой:

val ds = Seq(
  (0, 5), (0, 15), (1, 7), (1, 3),
  (2, 3), (2, 4), (2, 5), (2, 9), (3, 25)
).toDF("userID", "entID").as[(Int, Int)]
// ds: org.apache.spark.sql.Dataset[(Int, Int)] =[userID: int, entID: int]

val map = ds.rdd.groupByKey.collectAsMap
// map: scala.collection.Map[Int,Iterable[Int]] = Map(
//   2 -> CompactBuffer(3, 4, 5, 9), 1 -> CompactBuffer(7, 3),
//   3 -> CompactBuffer(25), 0 -> CompactBuffer(5, 15)
// )
Другие вопросы по тегам