Select the last element of an Array in a DataFrame - scala

Question

I'm working on a project and I'm dealing with some nested JSON date with a complicated schema/data structure. Basically what I want to do is filter out one of the columns, in a dataframe, such that I select the last element in the array. I'm totally stuck on how to do this. I hope this make sense.
Below is an example of what I'm trying to accomplish:
val singersDF = Seq(
("beatles", "help,hey,jude"),
("romeo", "eres,mia"),
("elvis", "this,is,an,example")
).toDF("name", "hit_songs")
val actualDF = singersDF.withColumn(
"hit_songs",
split(col("hit_songs"), "\\,")
)
actualDF.show(false)
actualDF.printSchema()
+-------+-----------------------+
|name |hit_songs |
+-------+-----------------------+
|beatles|[help, hey, jude] |
|romeo |[eres, mia] |
|elvis |[this, is, an, example]|
+-------+-----------------------+
root
|-- name: string (nullable = true)
|-- hit_songs: array (nullable = true)
| |-- element: string (containsNull = true)
The end goal for the output would be the following, to select the last "string" in the hit_songs array.
I'm not worried about what the schema would look like afterwards.
+-------+---------+
|name |hit_songs|
+-------+---------+
|beatles|jude |
|romeo |mia |
|elvis |example |
+-------+---------+

Answer 1

You can use the size function to calculate the index of the desired item in the array, and then pass this as the argument of Column.apply (explicitly or implicitly):
import org.apache.spark.sql.functions._
import spark.implicits._
actualDF.withColumn("hit_songs", $"hit_songs".apply(size($"hit_songs").minus(1)))
Or:
actualDF.withColumn("hit_songs", $"hit_songs"(size($"hit_songs").minus(1)))

Answer 2

Since spark 2.4+, you can use element_at which supports negative indexing. As you can see in this documentation quote:
element_at(array, index) - Returns element of array at given (1-based) index. If index < 0, accesses elements from the last to the first. Returns NULL if the index exceeds the length of the array.
With that, here's how to get the last element:
import org.apache.spark.sql.functions.element_at
actualDF.withColumn("hit_songs", element_at($"hit_songs", -1))
Reproducible example:
First let's prepare a sample dataframe with an array column:
val columns = Seq("col1")
val data = Seq((Array(1,2,3)))
val rdd = spark.sparkContext.parallelize(data)
val df = rdd.toDF(columns:_*)
which looks like this:
scala> df.show()
+---------+
| col1|
+---------+
|[1, 2, 3]|
+---------+
Then, apply element_at to get the last element as follows:
scala> df.withColumn("last_value", element_at($"col1", -1)).show()
+---------+----------+
| col1|last_value|
+---------+----------+
|[1, 2, 3]| 3|
+---------+----------+

Answer 3

Here's one approach:
val actualDF = Seq(
("beatles", Seq("help", "hey", "jude")),
("romeo", Seq("eres", "mia")),
("elvis", Seq("this", "is", "an", "example"))
).toDF("name", "hit_songs")
import org.apache.spark.sql.functions._
actualDF.withColumn("total_songs", size($"hit_songs")).
select($"name", $"hit_songs"($"total_songs" - 1).as("last_song"))
// +-------+---------+
// | name|last_song|
// +-------+---------+
// |beatles| jude|
// | romeo| mia|
// | elvis| example|
// +-------+---------+

Answer 4

You are looking for the SparkSQL function slice. or this PySpark Source.
Your implementation in Scala slice($"hit_songs", -1, 1)(0) where -1 is the starting position (last index) and 1 is the length, and (0) extracts the first string from resulting array of exactly 1 element.
Full Example:
import org.apache.spark.sql.functions._
val singersDF = Seq(
("beatles", "help,hey,jude"),
("romeo", "eres,mia"),
("elvis", "this,is,an,example")
).toDF("name", "hit_songs")
val actualDF = singersDF.withColumn(
"hit_songs",
split(col("hit_songs"), "\\,")
)
val newDF = actualDF.withColumn("last_song", slice($"hit_songs", -1, 1)(0))
display(newDF)
Output:
+---------+------------------------------+-------------+
| name | hit_songs | last_song |
+---------+------------------------------+-------------+
| beatles | ["help","hey","jude"] | jude |
| romeo | ["eres","mia"] | mia |
| elvis | ["this","is","an","example"] | example |
+---------+------------------------------+-------------+

Answer 5

You can also use an UDF like:
val lastElementUDF = udf((array: Seq[String]) => array.lastOption)
actualDF.withColumn("hit_songs", lastElementUDF($"hit_songs"))
array.lastOption will return None or Some, and array.last will throw an exception if the array is empty.