Exploding column with index

Question

I know that I can "explode" a column of type array like this:

import org.apache.spark.sql._
import org.apache.spark.sql.functions.explode
val explodedDf = 
    payloadLegsDf.withColumn("legs", explode(payloadLegsDf.col("legs")))

Now I have multiple rows; one for each item in the array.

Is there a way I can "explode with index"? So that there will be a new column that contains the index of the item in the original array?

(I can think of hacks to do this. First make the array field into an array of tuples of the original value and the index. Then do the explode. Then unpack the tuples. But is there a more elegant way?)

Answer 1

If you are using Spark 2.1+, the posexplode function can be used for that:

Creates a new row for each element with position in the given array or map column.

Example:

val df = Seq(
  (1L, Array[String]("a", "b")),
  (2L, Array[String]("c", "d"))
).toDF("id", "items")

val res = df.select($"id", posexplode($"items"))

This will create two new columns, pos for position/index and col for the extracted value:

+---+---+---+
| id|pos|col|
+---+---+---+
|  1|  0|  a|
|  1|  1|  b|
|  2|  0|  c|
|  2|  1|  d|
+---+---+---+