在 Spark Scala 中迭代数据框列 Array of Array答案

【问题标题】：Iterate dataframe column Array of Array in Spark Scala在 Spark Scala 中迭代数据框列 Array of Array
【发布时间】：2021-11-25 23:51:41
【问题描述】：

我正在尝试将数组数组作为 Spark 数据帧中的列进行迭代。寻找最好的方法来做到这一点。

架构：

root
 |-- Animal: struct (nullable = true)
 |    |-- Species: array (nullable = true)
 |    |    |-- element: struct (containsNull = true)
 |    |    |    |-- mammal: array (nullable = true)
 |    |    |    |    |-- element: struct (containsNull = true)
 |    |    |    |    |    |-- description: string (nullable = true)

目前我正在使用这个逻辑。这只会获取第一个数组。

df.select(
   col("Animal.Species").getItem(0).getItem("mammal").getItem("description")
)

伪逻辑：

col("Animal.Species").getItem(0).getItem("mammal").getItem("description")
+
col("Animal.Species").getItem(1).getItem("mammal").getItem("description")
+
col("Animal.Species").getItem(2).getItem("mammal").getItem("description")
+
col("Animal.Species").getItem(...).getItem("mammal").getItem("description")

所需的示例输出（扁平元素作为字符串）

llama, sheep, rabbit, hare

【问题讨论】：

标签： scala dataframe apache-spark apache-spark-sql

【解决方案1】：

不明显，但您可以使用.（或Column 的getField 方法）选择“通过”结构数组。选择Animal.Species.mammal 会返回一个由最内层结构组成的数组。不幸的是，这个数组数组会阻止您使用Animal.Species.mammal.description 之类的东西进一步向下钻取，因此您需要先将其展平，然后使用getField()。

如果我正确理解您的架构，则以下 JSON 应该是有效输入：

{
  "Animal": {
    "Species": [
      {
        "mammal": [
          { "description": "llama" },
          { "description": "sheep" }
        ]
      },
      {
        "mammal": [
          { "description": "rabbit" },
          { "description": "hare" }
        ]
      }
    ]
  }
}

val df = spark.read.json("data.json")
df.printSchema
// root
//  |-- Animal: struct (nullable = true)
//  |    |-- Species: array (nullable = true)
//  |    |    |-- element: struct (containsNull = true)
//  |    |    |    |-- mammal: array (nullable = true)
//  |    |    |    |    |-- element: struct (containsNull = true)
//  |    |    |    |    |    |-- description: string (nullable = true)

df.select("Animal.Species.mammal").show(false)
// +----------------------------------------+
// |mammal                                  |
// +----------------------------------------+
// |[[{llama}, {sheep}], [{rabbit}, {hare}]]|
// +----------------------------------------+

df.select(flatten(col("Animal.Species.mammal"))).show(false)
// +------------------------------------+
// |flatten(Animal.Species.mammal)      |
// +------------------------------------+
// |[{llama}, {sheep}, {rabbit}, {hare}]|
// +------------------------------------+

现在这是一个结构数组，您可以使用getField("description") 获取感兴趣的数组：

df.select(flatten(col("Animal.Species.mammal")).getField("description")).show(false)
// +--------------------------------------------------------+
// |flatten(Animal.Species.mammal AS mammal#173).description|
// +--------------------------------------------------------+
// |[llama, sheep, rabbit, hare]                            |
// +--------------------------------------------------------+

最后可以使用array_join加上分隔符", "来获取想要的字符串：

df.select(
  array_join(
    flatten(col("Animal.Species.mammal")).getField("description"),
    ", "
  ) as "animals"
).show(false)
// +--------------------------+
// |animals                   |
// +--------------------------+
// |llama, sheep, rabbit, hare|
// +--------------------------+

【讨论】：

这是一个优雅而简单的解决方案！太棒了！

【解决方案2】：

您可以申请两次explode：第一次申请Animal.Species，第二次申请第一次的结果：

import org.apache.spark.sql.functions._
df.withColumn("tmp", explode(col("Animal.Species")))
  .withColumn("tmp", explode(col("tmp.mammal")))
  .select("tmp.description")
  .show()

【讨论】：