Spark: индекс значения в массиве внутри ColumnType

Question

В Spark, используя Scala, у меня есть набор данных со столбцом, имеющим массив, как

[foo, bar, baz, bei]
[foo, bar, baz, bei]
[foo, zie]

Теперь я хочу добавить еще один столбец с индексом значения "бар"

[foo, bar, baz, bei]     2  
[foo, zie]  0 

Есть ли что-то похожее на

.withColumn("idx", array_contains(col("Name"),"bar"))

который возвращает мне истину / ложь (мне нужен индекс значения)?

Answer 1

С UDF:

val df = List(
  Seq("foo", "bar", "baz", "bei"),
  Seq("foo", "bar", "baz", "bei"),
  Seq("foo", "zie")
).toDF()

val getIndex = (seq: Seq[String]) => seq.indexOf("bar") + 1
val getIndexUDF = udf(getIndex)

val result = df.withColumn("idx", getIndexUDF($"value"))
result.show(false)

Выход:

+--------------------+---+
|value               |idx|
+--------------------+---+
|[foo, bar, baz, bei]|2  |
|[foo, bar, baz, bei]|2  |
|[foo, zie]          |0  |
+--------------------+---+

Answer 2

Начиная с версии 2.4.0 Spark предоставляет array_position функция

import org.apache.spark.sql.functions.array_position

df.withColumn("idx", array_position($"Name", "bar"))