使用Scala从Spark中的嵌套结构数据帧获取值

时间:2018-08-28 10:53:11

标签: scala apache-spark dataframe

我有一个具有嵌套结构(数组数组)的数据框,

StructField("Games", ArrayType(StructType(Array(
    StructField("Team", StringType, true),
    StructField("Amount", StringType, true),
    StructField("Game", StringType, true)))), true),

为此,我将获得如下所示的值(团队,数量,游戏按照此处的顺序)

[[A,160,Chess], [B,100,Hockey], [C,1200,Football], [D,900,Cricket]]
[[E,700,Cricket], [F,1000,Chess]]
[[G,1900,Basketball], [I,1000,Cricket], [H,9000,Football]]

现在,如果
,我必须从此数据框中获取值 Game === 'Football' then TeamFootball = C and Amount = 1200
Game === 'Cricket' then TeamCricket = D and Amount = 900用于第一行。

我尝试过

.withColumn("TeamFootball", when($"Games.Game".getItem(2)==="Football",$"Games.Team".getItem(0).cast(StringType)).otherwise(lit("NA")))
.withColumn("TeamCricket", when($"Games.Game".getItem(2)==="Cricket", $"Games.Team".getItem(0).cast(StringType)).otherwise(lit("NA")))
.withColumn("TeamFootballAmount", when($"Games.Game".getItem(2)==="Football",$"Games.Amount".getItem(1).cast(StringType)).otherwise(lit("NA")))
.withColumn("TeamCricketAmount", when($"Games.Game".getItem(2)==="Cricket",$"Games.Amount".getItem(1).cast(StringType)).otherwise(lit("NA")))

我需要同一行中的所有这些列,这就是为什么我不使用explode的原因。 在这里,我无法处理数组索引,请您帮忙。

1 个答案:

答案 0 :(得分:1)

先“分解”然后再进行“透视”会有所帮助,请在输出中检查“结果”:

val data = List(
  (1, "A", 160, "Chess"), (1, "B", 100, "Hockey"), (1, "C", 1200, "Football"), (1, "D", 900, "Cricket"),
  (2, "E", 700, "Cricket"), (2, "F", 1000, "Chess"),
  (3, "G", 1900, "Basketball"), (3, "I", 1000, "Cricket"), (3, "H", 9000, "Football")
)
val unstructured = data.toDF("id", "Team", "Amount", "Game")
unstructured.show(false)

val original = unstructured.groupBy("id").agg(collect_list(struct($"Team", $"Amount", $"Game")).alias("Games"))
println("--- Original ----")
original.printSchema()
original.show(false)

val exploded = original.withColumn("Games", explode($"Games")).select("id", "Games.*")
println("--- Exploded ----")
exploded.show(false)

println("--- Result ----")
exploded.groupBy("id").pivot("Game").agg(max($"Amount").alias("Amount"), max("Team").alias("Team")).orderBy("id").show(false)

输出为:

+---+----+------+----------+
|id |Team|Amount|Game      |
+---+----+------+----------+
|1  |A   |160   |Chess     |
|1  |B   |100   |Hockey    |
|1  |C   |1200  |Football  |
|1  |D   |900   |Cricket   |
|2  |E   |700   |Cricket   |
|2  |F   |1000  |Chess     |
|3  |G   |1900  |Basketball|
|3  |I   |1000  |Cricket   |
|3  |H   |9000  |Football  |
+---+----+------+----------+

--- Original ----
root
 |-- id: integer (nullable = false)
 |-- Games: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- Team: string (nullable = true)
 |    |    |-- Amount: integer (nullable = false)
 |    |    |-- Game: string (nullable = true)

+---+-------------------------------------------------------------------+
|id |Games                                                              |
+---+-------------------------------------------------------------------+
|3  |[[G,1900,Basketball], [I,1000,Cricket], [H,9000,Football]]         |
|1  |[[A,160,Chess], [B,100,Hockey], [C,1200,Football], [D,900,Cricket]]|
|2  |[[E,700,Cricket], [F,1000,Chess]]                                  |
+---+-------------------------------------------------------------------+

--- Exploded ----
+---+----+------+----------+
|id |Team|Amount|Game      |
+---+----+------+----------+
|3  |G   |1900  |Basketball|
|3  |I   |1000  |Cricket   |
|3  |H   |9000  |Football  |
|1  |A   |160   |Chess     |
|1  |B   |100   |Hockey    |
|1  |C   |1200  |Football  |
|1  |D   |900   |Cricket   |
|2  |E   |700   |Cricket   |
|2  |F   |1000  |Chess     |
+---+----+------+----------+

--- Result ----
+---+-----------------+---------------+------------+----------+--------------+------------+---------------+-------------+-------------+-----------+
|id |Basketball_Amount|Basketball_Team|Chess_Amount|Chess_Team|Cricket_Amount|Cricket_Team|Football_Amount|Football_Team|Hockey_Amount|Hockey_Team|
+---+-----------------+---------------+------------+----------+--------------+------------+---------------+-------------+-------------+-----------+
|1  |null             |null           |160         |A         |900           |D           |1200           |C            |100          |B          |
|2  |null             |null           |1000        |F         |700           |E           |null           |null         |null         |null       |
|3  |1900             |G              |null        |null      |1000          |I           |9000           |H            |null         |null       |
+---+-----------------+---------------+------------+----------+--------------+------------+---------------+-------------+-------------+-----------+