从 Spark 中的元素获取 Seq 索引 Sql

Question

我有一个如下所示的数据框：

---------------------+------------------------
text                 | featured_text
---------------------+------------------------
sun                  | [type, move, sun]
---------------------+------------------------

我想搜索 "featured_text" 数组中的 "text" 列值并获取 "text" 值的索引（如果存在）。在上面的示例中，我想在数组 [type, move, sun] 中搜索 "sun"，结果将是“2”（索引）。

是否有任何 spark sql function/scala 函数可用于从元素中获取索引？

Answer 1

据我所知，没有函数可以直接使用 Spark SQL API 执行此操作。但是，您可以使用 UDF 代替（我假设输入数据帧称为 df）：

val getIndex = udf((text: String, featuredText: Seq[String]) => {
  featuredText.indexOf(text)
})

val df2 = df.withColumn("index", getIndex($"text", $"featured_text"))

这将给出：

+----+-----------------+-----+
|text|    featured_text|index|
+----+-----------------+-----+
| sun|[type, move, sun]|    2|
+----+-----------------+-----+

如果值不存在，索引列将具有 -1。