将结构字段与另一列 Spark 进行比较

Question

我有一个带有 ArrayType 结构字段的数据集，我需要过滤与 max_stat 值一致的 stat 字段的值。数据结构如下：

 |-- data: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- id: integer (nullable = true)
 |    |    |-- stat: float (nullable = true)
 |-- max_stat: float (nullable = true)
 |-- naming: string (nullable = true)

我试图做这样的事情，但无法将 Array 与 int 进行比较。那么如何更改此过滤器？

df.filter($"data.stat" === "max_stat")
.drop(max_stat)

输入数据：

 data                           | max_stat | naming
[(1,0.34),(2, 0.57), (3, 0.89)] | 0.89     | example

预期输出：

 data     | naming
(3, 0.89) | example

Answer 1

您必须使用explode 才能为data 中的每个element 生成一行。

import org.apache.spark.sql.functions.explode

ds
  .withColumn("data", explode($"data"))
  .filter($"data.stat" === $"max_stat")
  .drop($"max_stat")
  .show()

输出：

+---------+-------+
|     data| naming|
+---------+-------+
|[3, 0.89]|example|
+---------+-------+

但是，explode 是一项非常昂贵的操作，如果您的数据集很大，可能会成为一个问题。另一种不使用explode的方法是：

import org.apache.spark.sql.functions._

ds
  .filter(array_contains($"data.stat", $"max_stat"))
  .withColumn("max_stat_idx", array_position($"data.stat", $"max_stat").cast(IntegerType))
  .withColumn("data", element_at($"data", $"max_stat_idx"))
  .drop("max_stat", "max_stat_idx")

基本上就是在data数组中寻找匹配值的索引，然后利用这个索引得到正确的element