Element position from nested DataFrame array (Spark 2.2) - scala

Question

I'm trying to explode a nested DataFrame in Spark Scala. I have a DataFrame df which contains the following information:
root
|-- id: integer (nullable = false)
|-- features: array (nullable = true)
| |-- element: float (containsNull = false)
I've exploded the array information into to a flat DataFrame with:
df.selectExpr("id","explode(features) as features")
and got the following DataFrame:
id features
0 0.0629885
0 0.15931357
0 0.08922347
My end goal is to pivot the data and calculate some similarities with that information. To do that, it would be very cool to get the actual position of the feature for every ID into the DataFrame, like this:
id features feature_pos
0 0.0629885 0
0 0.15931357 1
0 0.08922347 2

Answer 1

Use posexplode in place of explode:
Creates a new row for each element with position in the given array or map column.
Unlike posexplode, if the array/map is null or empty then the row (null, null) is produced.

Answer 2

Here is the example with posexplode.
scala> val df = Seq((0, Seq(0.1f, 0.2f, 0.3f)),(1, Seq(0.4f, 0.5f, 0.6f))).toDF("id", "features")
df: org.apache.spark.sql.DataFrame = [id: int, features: array<float>]
scala> df.show(false)
+---+---------------+
|id |features |
+---+---------------+
|0 |[0.1, 0.2, 0.3]|
|1 |[0.4, 0.5, 0.6]|
+---+---------------+
Note that df.withColumn("pos cols",posexplode('features)).show(false) will throw error, so use df.select()
scala> df.select(posexplode('features)).show(false)
+---+---+
|pos|col|
+---+---+
|0 |0.1|
|1 |0.2|
|2 |0.3|
|0 |0.4|
|1 |0.5|
|2 |0.6|
+---+---+
scala>
The default names are "pos" and "col". You can rename them as
scala> df.select(posexplode('features).as(Seq("a","b"))).show(false)
+---+---+
|a |b |
+---+---+
|0 |0.1|
|1 |0.2|
|2 |0.3|
|0 |0.4|
|1 |0.5|
|2 |0.6|
+---+---+
scala>
When you want to explode and select all columns, use
scala> df.select(col("*"), posexplode('features).as( Seq("a","b")) ).show(false)
+---+---------------+---+---+
|id |features |a |b |
+---+---------------+---+---+
|0 |[0.1, 0.2, 0.3]|0 |0.1|
|0 |[0.1, 0.2, 0.3]|1 |0.2|
|0 |[0.1, 0.2, 0.3]|2 |0.3|
|1 |[0.4, 0.5, 0.6]|0 |0.4|
|1 |[0.4, 0.5, 0.6]|1 |0.5|
|1 |[0.4, 0.5, 0.6]|2 |0.6|
+---+---------------+---+---+
scala>

Answer 3

You can also apply Scala's zipWithIndex via a UDF as follows:
val df = Seq(
(0, Seq(0.1f, 0.2f, 0.3f)),
(1, Seq(0.4f, 0.5f, 0.6f))
).toDF("id", "features")
def addIndex = udf(
(s: Seq[Float]) => s.zipWithIndex
)
val df2 = df.withColumn( "features_idx", explode(addIndex($"features")) )
df2.select( $"id", $"features_idx._1".as("features"), $"features_idx._2".as("features_pos") ).show
+---+--------+------------+
| id|features|features_pos|
+---+--------+------------+
| 0| 0.1| 0|
| 0| 0.2| 1|
| 0| 0.3| 2|
| 1| 0.4| 0|
| 1| 0.5| 1|
| 1| 0.6| 2|
+---+--------+------------+