根据给定条件从 array<struct> 中选择一行

Question

我有一个具有以下架构的数据框 -

|-- ID: string (nullable = true)
|-- VALUES: array (nullable = true)
|    |-- element: struct (containsNull = true)
|    |    |-- _v1: string (nullable = true)
|    |    |-- _v2: string (nullable = true)

值就像 -

[["ABC","a"],["PQR","c"],["XYZ","b"],["TUV","d"]]
[["PQR","g"]]
[["TUV","f"],["ABC","e"]]

我必须根据 _v1 的值从这个数组中选择一个结构。这些值中有一个层次结构，例如 -

“ABC”->“XYZ”->“PQR”->“TUV”

现在，如果存在“TUV”，我们将选择_v1 中带有“TUV”的行。否则我们将检查“PQR”。如果存在“PQR”，则取其所在的行。否则检查“XYZ”等等。

结果 df 应该看起来像 - （现在将是 StructType，而不是 Array[Struct]）

["TUV","d"]
["PQR","g"]
["TUV","f"]

有人可以指导我如何通过创建 udf 来解决这个问题吗？ 提前致谢。

Answer 1

您可以执行以下操作

import org.apache.spark.sql.functions._
def getIndex = udf((array : mutable.WrappedArray[String]) => {
  if(array.contains("TUV")) array.indexOf("TUV")
  else if(array.contains("PQR")) array.indexOf("PQR")
  else if(array.contains("XYZ")) array.indexOf("XYZ")
  else if(array.contains("ABC")) array.indexOf("ABC")
  else 0
})

df.select($"VALUES"(getIndex($"VALUES._v1")).as("selected"))

你应该有以下输出

+--------+
|selected|
+--------+
|[TUV,d] |
|[PQR,g] |
|[TUV,f] |
+--------+

希望回答对你有帮助

更新

您可以使用. 表示法选择struct 列的元素。这里$"VALUES._v1" 选择struct 中的所有_v1，并将它们以Array 的顺序传递给udf 函数。

例如：对于[["ABC","a"],["PQR","c"],["XYZ","b"],["TUV","d"]]，$"VALUES._v1" 将返回["ABC","PQR","XYZ","TUV"]，并传递给udf 函数

在udf函数内部，返回匹配字符串的数组索引。例如：对于 ["ABC","PQR","XYZ","TUV"]，"TUV" 匹配，所以它会返回 3。

对于第一行，getIndex($"VALUES._v1") 将返回 3，因此 $"VALUES"(getIndex($"VALUES._v1") 等价于 $"VALUES"(3)，这是 [["ABC","a"],["PQR","c"],["XYZ","b"],["TUV","d"]] 的第四个元素，即 ["TUV","d"]。

我希望解释清楚。

Answer 2

只要每行最多只包含每个_v1 值一次，这应该可以工作。 UDF 将返回 hierarchy 列表中最佳值的索引。然后_v1中包含该值的stuct将被选中并放入“select”列。

val hierarchy = List("TUV", "PQR", "XYZ", "ABC")

val findIndex = udf((rows: Seq[String]) => {
  val s = rows.toSet
  val best = hierarchy.filter(h => s contains h).head
  rows.indexOf(best)
})

df.withColumn("select", $"VALUES"(findIndex($"VALUES._v2")))

一个列表用于订单，以便轻松扩展到超过 4 个值。