Counting longest sequence of specific elements in a list contained within a spark.sql database column

Viewed 196

I have the following problem I would like to solve.

I have the following Dataframe that I created from a query

val temp = spark.sql("select Id, collect_list(from) as letter from f group by Id")

|Id|                letter|
+-----------+---------------+
|      106|            [c]|
|      101|            [p]|
|      104|[c, c, c, t, u]|
|      100|[d, t, j, j, c]|
|      110|      [p, n, f]|
|      113|[s, c, c, b, ..|
|      115|[u, s, t, c, ..|
|       11|   [c, c, i, s]|
|      117|   [d, d, p, s]|
|      118|[a, s, c, t, ..|
|      123|         [d, n]|
|      125|         [n, b]|
|      128|            [c]|
|      131|   [c, t, c, u]|
|      132|      [c, u, i]|
|      134|[c, p, j, u, c]|
|      136|[b, a, t, n, c]|
|      137|         [b, a]|
|      138|      [b, t, c]|
|      141|            [s]| 

I would like to create a new column called "n" This column would contain a numerical value which represents the longest sequence of letters in a cell before "c" appears. the longest sequence can be anywhere in the list.

For example the solution column for this section (assuming nothing is cut off by the ....) would be

0, 1, 3, 5, 3, 2, 4, 4, 4, 4, 2, 2, 1, 4, 2, 5, 5, 2, 3, 1

Any help would be greatly appreciated. Thank you!

2 Answers

You could write a user defined function (udf) that would compute what you wish. There are plenty of ways to compute the longuest sequence. One simple way is to split the sequence on "co", compute the size of each sub sequence and take the max.

val longuest_seq = udf((x : Seq[String]) => {
    x.reduce(_ +" "+_)
     .split(" *co *")
     .map(_.count(_ == ' ') + 1)
     .max
})

val df = Seq(
    (1, Array("x", "y", "co", "z")),
    (2, Array("x")),
    (3, Array("co", "t")),
    (4, Array("a", "b", "c", "d", "co", "e"))
).toDF("id", "trip")

df.withColumn("n_trips", longuest_seq('trip)).show

which yields

+---+-------------------+-------+
| id|               trip|n_trips|
+---+-------------------+-------+
|  1|      [x, y, co, z]|      2|
|  2|                [x]|      1|
|  3|            [co, t]|      1|
|  4|[a, b, c, d, co, e]|      4|
+---+-------------------+-------+

Here is how you can use the spark functions, you can convert the given scala functions with spark functions as below

import org.apache.spark.sql.functions._

df.withColumn("n_trip",
  array_max(
    transform(
      filter(
        split(array_join($"trip", " "), "co"),
        (col: Column) => (col =!= "" || col =!= null)
      ), (col: Column) => size(split(trim(col), " "))
    )
  ))
  .withColumn("n_trip", when($"n_trip".isNull, 0).otherwise($"n_trip"))
  .show(false)

Update: Easy to understand

  df.withColumn("split", split(array_join($"trip", " "), "co"))
    .withColumn("filter", filter($"split", (col: Column) => col =!= "" || col =!= null))
    .withColumn("n_trip", array_max(transform($"filter", (col: Column) => size(split(trim(col), " ")))))
    .withColumn("n_trip", when($"n_trip".isNull, 0).otherwise($"n_trip"))
    .drop("split", "filter")
    .show(false)

Output:

+-----------+--------------------+------+
|passengerId|trip                |n_trip|
+-----------+--------------------+------+
|10096      |[co]                |0     |
|10351      |[pk]                |1     |
|10436      |[co, co, cn, tj, us]|3     |
|1090       |[dk, tj, jo, jo, ch]|5     |
|11078      |[pk, no, fr]        |3     |
|11332      |[sg, cn, co, bm]    |2     |
|11563      |[us, sg, th, cn]    |4     |
|1159       |[ca, cl, il, sg]    |4     |
|11722      |[dk, dk, pk, sg]    |4     |
|11888      |[au, se, ca, tj]    |4     |
|12394      |[dk, nl]            |2     |
|12529      |[no, be]            |2     |
|12847      |[cn]                |1     |
|13192      |[cn, tk, cg, uk]    |4     |
|13282      |[co, us, iq]        |2     |
|13442      |[cn, pk, jo, us, ch]|5     |
|13610      |[be, ar, tj, no, ch]|5     |
|13772      |[be, at]            |2     |
|13865      |[be, th, cn]        |3     |
|14157      |[sg]                |1     |
+-----------+--------------------+------+
Related