Question

在Spark中，使用Scala，我有一个数据集，其中的列具有类似数组的

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

现在我想添加索引值为“ bar”的另一列

[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"))

Spark：ColumnType中数组中值的索引

2 个答案: