Apache Spark - Scala API - Aggregate on sequentially increasing key

Question

I have a data frame that looks something like this:

val df = sc.parallelize(Seq(
  (3,1,"A"),(3,2,"B"),(3,3,"C"),
  (2,1,"D"),(2,2,"E"),
  (3,1,"F"),(3,2,"G"),(3,3,"G"),
  (2,1,"X"),(2,2,"X")
)).toDF("TotalN", "N", "String")

+------+---+------+
|TotalN|  N|String|
+------+---+------+
|     3|  1|     A|
|     3|  2|     B|
|     3|  3|     C|
|     2|  1|     D|
|     2|  2|     E|
|     3|  1|     F|
|     3|  2|     G|
|     3|  3|     G|
|     2|  1|     X|
|     2|  2|     X|
+------+---+------+

I need to aggregate the strings by concatenating them together based on the TotalN and the sequentially increasing ID (N). The problem is there is not a unique ID for each aggregation I can group by. So, I need to do something like "for each row look at the TotalN, loop through the next N rows and concatenate, then reset".

+------+------+
|TotalN|String|
+------+------+
|     3|   ABC|
|     2|    DE|
|     3|   FGG|
|     2|    XX|
+------+------+

Any pointers much appreciated.

Using Spark 2.3.1 and the Scala Api.

Answer 1

Try this:

val df = spark.sparkContext.parallelize(Seq(
  (3, 1, "A"), (3, 2, "B"), (3, 3, "C"),
  (2, 1, "D"), (2, 2, "E"),
  (3, 1, "F"), (3, 2, "G"), (3, 3, "G"),
  (2, 1, "X"), (2, 2, "X")
)).toDF("TotalN", "N", "String")


df.createOrReplaceTempView("data")

val sqlDF = spark.sql(
  """
    | SELECT TotalN d, N, String, ROW_NUMBER() over (order by TotalN) as rowNum
    | FROM data
  """.stripMargin)

sqlDF.withColumn("key", $"N" - $"rowNum")
  .groupBy("key").agg(collect_list('String).as("texts")).show()

Answer 2

Solution is to calculate a grouping variable using the row_number function which can be used in later groupBy.

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions.row_number

var w = Window.orderBy("TotalN")
df.withColumn("GeneratedID", $"N" - row_number.over(w)).show

+------+---+------+-----------+
|TotalN|  N|String|GeneratedID|
+------+---+------+-----------+
|     2|  1|     D|          0|
|     2|  2|     E|          0|
|     2|  1|     X|         -2|
|     2|  2|     X|         -2|
|     3|  1|     A|         -4|
|     3|  2|     B|         -4|
|     3|  3|     C|         -4|
|     3|  1|     F|         -7|
|     3|  2|     G|         -7|
|     3|  3|     G|         -7|
+------+---+------+-----------+